| import importlib | |
| import io | |
| import os | |
| from pathlib import Path | |
| from tempfile import TemporaryDirectory | |
| from django.test import TestCase | |
| from hue_portal.core.models import Office, Fine | |
| class EtlLoaderTestCase(TestCase): | |
| def setUp(self): | |
| self.tempdir = TemporaryDirectory() | |
| self.data_dir = Path(self.tempdir.name) | |
| self._write_office_csv() | |
| self._write_fine_csv() | |
| def tearDown(self): | |
| self.tempdir.cleanup() | |
| os.environ.pop("ETL_DATA_DIR", None) | |
| def _write_office_csv(self): | |
| path = self.data_dir / "danh_ba_diem_tiep_dan.csv" | |
| path.write_text( | |
| "unit_name,address,district,working_hours,phone,email,latitude,longitude,service_scope,updated_at\n" | |
| "Công an phường A,123 Đường B,Quận 1,08:00-17:00,0123456789,ca@example.com,16.0,108.0,Tiếp dân,2025-01-01\n", | |
| encoding="utf-8" | |
| ) | |
| def _write_fine_csv(self): | |
| path = self.data_dir / "muc_phat_theo_hanh_vi.csv" | |
| path.write_text( | |
| "violation_code,violation_name,article,decree,min_fine,max_fine,license_points,remedial_measures,source_url,updated_at\n" | |
| "V001,Vượt đèn đỏ,5,100/2019/NĐ-CP,1000000,3000000,2,Phạt bổ sung,http://example.com,2025-01-01\n", | |
| encoding="utf-8" | |
| ) | |
| def _load_module(self): | |
| os.environ["ETL_DATA_DIR"] = str(self.data_dir) | |
| module = importlib.import_module("scripts.etl_load") | |
| return importlib.reload(module) | |
| def test_load_offices_creates_records(self): | |
| etl = self._load_module() | |
| log_buffer = io.StringIO() | |
| processed = etl.load_offices(since=None, dry_run=False, log_file=log_buffer) | |
| self.assertEqual(processed, 1) | |
| self.assertEqual(Office.objects.count(), 1) | |
| office = Office.objects.first() | |
| self.assertEqual(office.unit_name, "Công an phường A") | |
| def test_load_fines_creates_records(self): | |
| etl = self._load_module() | |
| log_buffer = io.StringIO() | |
| processed = etl.load_fines(since=None, dry_run=False, log_file=log_buffer) | |
| self.assertEqual(processed, 1) | |
| self.assertEqual(Fine.objects.count(), 1) | |
| fine = Fine.objects.first() | |
| self.assertEqual(fine.code, "V001") | |
| if __name__ == "__main__": | |
| import unittest | |
| unittest.main() | |