from datetime import datetime

import pytest
from model_mommy import mommy

from .models import Filiacao, Legislatura, Mandato

pytestmark = pytest.mark.django_db


def data(valor):
    return datetime.strptime(valor, '%Y-%m-%d').date()


def test_filiacoes():
    legislatura = mommy.make(Legislatura,
                             data_inicio=data('2001-01-01'),
                             data_fim=data('2001-12-31'),
                             )
    mandato = mommy.make(Mandato, legislatura=legislatura)
    f1_fora, f2, f3, f4 = [mommy.make(Filiacao,
                                      parlamentar=mandato.parlamentar,
                                      data=ini,
                                      data_desfiliacao=fim)
                           for ini, fim in (
        (data('2000-01-01'), data('2000-12-31')),
        (data('2000-01-01'), data('2001-03-01')),
        (data('2001-03-02'), data('2001-10-01')),
        (data('2001-10-02'), None),
    )]
    assert mandato.get_partidos() == [f2.partido, f3.partido, f4.partido]