diff --git a/jarbas/core/management/commands/companies.py b/jarbas/core/management/commands/companies.py index 451c4c8..5ad7dbb 100644 --- a/jarbas/core/management/commands/companies.py +++ b/jarbas/core/management/commands/companies.py @@ -1,13 +1,24 @@ -import csv import lzma - -from django.core.exceptions import ValidationError -from django.core.validators import validate_email +import rows from jarbas.core.management.commands import LoadCommand from jarbas.core.models import Activity, Company +class CompaniesDate(rows.fields.DateField): + INPUT_FORMAT = '%d/%m/%Y' + + +COMPANIES_CSV_FIELD_TYPES = { + 'email': rows.fields.EmailField, + 'opening': CompaniesDate, + 'situation_date': CompaniesDate, + 'special_situation_date': CompaniesDate, + 'latitude': rows.fields.FloatField, + 'longitude': rows.fields.FloatField +} + + class Command(LoadCommand): help = 'Load Serenata de Amor companies dataset into the database' @@ -30,12 +41,14 @@ def save_companies(self): """ skip = ('main_activity', 'secondary_activity') keys = tuple(f.name for f in Company._meta.fields if f not in skip) - with lzma.open(self.path, mode='rt', encoding='utf-8') as file_handler: - for row in csv.DictReader(file_handler): + with lzma.open(self.path, mode='rb') as file_handler: + for row in rows.import_from_csv(file_handler, force_types=COMPANIES_CSV_FIELD_TYPES, encoding='utf-8'): + row = dict(row._asdict()) + main, secondary = self.save_activities(row) filtered = {k: v for k, v in row.items() if k in keys} - obj = Company.objects.create(**self.serialize(filtered)) + obj = Company.objects.create(**filtered) for activity in main: obj.main_activity.add(activity) for activity in secondary: @@ -62,25 +75,3 @@ def save_activities(self, row): secondaries.append(obj) return [main], secondaries - - def serialize(self, row): - row['email'] = self.to_email(row['email']) - - dates = ('opening', 'situation_date', 'special_situation_date') - for key in dates: - row[key] = self.to_date(row[key]) - - decimals = ('latitude', 'longitude') - for key in decimals: - row[key] = self.to_number(row[key]) - - return row - - @staticmethod - def to_email(email): - try: - validate_email(email) - return email - - except ValidationError: - return None diff --git a/jarbas/core/tests/fixtures/companies.xz b/jarbas/core/tests/fixtures/companies.xz new file mode 100644 index 0000000..09025f3 Binary files /dev/null and b/jarbas/core/tests/fixtures/companies.xz differ diff --git a/jarbas/core/tests/test_companies_command.py b/jarbas/core/tests/test_companies_command.py index f9995d6..1996da2 100644 --- a/jarbas/core/tests/test_companies_command.py +++ b/jarbas/core/tests/test_companies_command.py @@ -1,12 +1,12 @@ -from datetime import date -from io import StringIO +import datetime +import os from unittest.mock import patch +from django.conf import settings from django.test import TestCase from jarbas.core.management.commands.companies import Command from jarbas.core.models import Activity, Company -from jarbas.core.tests import sample_company_data class TestCommand(TestCase): @@ -15,33 +15,6 @@ def setUp(self): self.command = Command() -class TestSerializer(TestCommand): - - def test_to_email(self): - expected = 'jane@example.com' - self.assertEqual(self.command.to_email('abc'), None) - self.assertEqual(self.command.to_email('jane@example.com'), expected) - - def test_serializer(self): - company = { - 'email': 'ahoy', - 'opening': '31/12/1969', - 'situation_date': '31/12/1969', - 'special_situation_date': '31/12/1969', - 'latitude': '3.1415', - 'longitude': '-42' - } - expected = { - 'email': None, - 'opening': date(1969, 12, 31), - 'situation_date': date(1969, 12, 31), - 'special_situation_date': date(1969, 12, 31), - 'latitude': 3.1415, - 'longitude': -42.0 - } - self.assertEqual(self.command.serialize(company), expected) - - class TestCreate(TestCommand): @patch.object(Activity.objects, 'update_or_create') @@ -61,22 +34,45 @@ def test_save_activities(self, update_or_create): self.assertEqual(1, len(main)) self.assertEqual(99, len(secondaries)) - @patch('jarbas.core.management.commands.companies.lzma') - @patch('jarbas.core.management.commands.companies.csv.DictReader') @patch('jarbas.core.management.commands.companies.Command.save_activities') - @patch('jarbas.core.management.commands.companies.Command.serialize') @patch('jarbas.core.management.commands.companies.Command.print_count') @patch.object(Company.objects, 'create') - def test_save_companies(self, create, print_count, serialize, save_activities, rows, lzma): + def test_save_companies(self, create, print_count, save_activities): self.command.count = 0 - lzma.return_value = StringIO() - rows.return_value = [sample_company_data] - serialize.return_value = dict(ahoy=42) save_activities.return_value = ([3], [14, 15]) - self.command.path = 'companies.xz' + self.command.path = os.path.join(settings.BASE_DIR, 'jarbas', 'core', 'tests', + 'fixtures', 'companies.xz') self.command.save_companies() - create.assert_called_with(ahoy=42) + expected = { + 'additional_address_details': b'', + 'address': b'', + 'city': b'', + 'cnpj': b'', + 'email': 'test@test.com.br', + 'last_updated': b'', + 'latitude': -15.7910966, + 'legal_entity': b'', + 'longitude': -47.9508743, + 'name': 'Test', + 'neighborhood': b'', + 'number': 1, + 'opening': None, + 'phone': b'', + 'responsible_federative_entity': b'', + 'situation': b'', + 'zip_code': b'', + 'situation_date': datetime.date(2005, 9, 24), + 'situation_reason': b'', + 'type': 'Book', + 'special_situation_date': None, + 'state': b'', + 'status': b'', + 'trade_name': b'', + 'special_situation': b'' + } + create.assert_called_with(**expected) create.return_value.main_activity.add.assert_called_with(3) + self.assertEqual(1, print_count.call_count) self.assertEqual(2, create.return_value.secondary_activity.add.call_count)