import pytest

from company.models.company import Company
from company.models.location import LocationCode
from company.services.company_create_update import CompanyCreateAndUpdateService
from ncircc.enums.notifications import AffectedSystemFunction


@pytest.mark.django_db
class TestCreateOrUpdateCompany:
    @pytest.fixture(autouse=True)
    def setup_test(self):
        self.location, _ = LocationCode.objects.get_or_create(code='RU-MOS')
        self.data_for_create = {
            'name': 'TestNAMEcompany',
            'is_cii': False,
            'location': self.location.id,
            'city': 'Moscow',
            'affected_system_function': AffectedSystemFunction.OTHER.value,
            'api_key': ''
        }

    @pytest.mark.unit
    def test_create_company(self):
        count_before = Company.objects.count()
        assert count_before == 0
        data = CompanyCreateAndUpdateService(None, self.data_for_create).save()

        assert data['name'] == 'TestNAMEcompany'
        assert not data['is_cii']
        assert data['location'] == {'id': self.location.id, 'code': self.location.code}
        assert data['city'] == 'Moscow'
        assert data['affected_system_function'] == AffectedSystemFunction.OTHER.value
        assert data['api_key'] == ''

        count_after = Company.objects.count()
        assert count_after == 1

    @pytest.mark.unit
    def test_update_partially(self):
        data_for_create = self.data_for_create.copy()
        data_for_create['location_id'] = self.location.id
        del data_for_create['location']
        company = Company.objects.create(**data_for_create)
        count_before = Company.objects.count()
        assert count_before == 1
        data_for_update = {
            'name': 'TestNameCompany2',
            'is_cii': True,
            'affected_system_function': AffectedSystemFunction.NUCLEAR_POWER.value,
        }
        data = CompanyCreateAndUpdateService(company, data_for_update).save()
        count_after = Company.objects.count()
        assert count_after == count_before
        assert data.get('name') == data_for_update.get('name')
        assert data.get('is_cii') == data_for_update.get('is_cii')
        assert data.get('affected_system_function') == 'Атомная энергетика'