"""Tests for organization source backfill tasks and schedules.""" from importlib import import_module from tempfile import TemporaryDirectory from apps.parsers.models import ParserLoadLog from django.apps import apps as django_apps from django.conf import settings from django.core.cache import cache from django.test import TestCase, override_settings from django.utils import timezone from django_celery_beat.models import PeriodicTask from organizations.cache import get_organization_api_cache_version from organizations.models import ( IndustrialProductionExtension, Organization, OrganizationSourceRecord, ) from organizations.tasks import ( backfill_all_organization_sources, backfill_organization_sources_for_parser_batch, refresh_source_record_export_artifacts, ) from tests.apps.parsers.factories import IndustrialCertificateRecordFactory class OrganizationSourceBackfillTasksTest(TestCase): """Checks Celery tasks that maintain API v2 organization source extensions.""" @staticmethod def _directory_organization(**kwargs) -> Organization: kwargs.setdefault("directory_imported_at", timezone.now()) return Organization.objects.create(**kwargs) def test_backfill_all_task_rebuilds_sources_and_invalidates_api_cache(self): organization = self._directory_organization( name='ООО "Источник"', inn="7800000401", ogrn="1027700144401", ) IndustrialCertificateRecordFactory( inn=organization.inn, ogrn=organization.ogrn, certificate_number="FULL-SOURCE-CERT", ) cache.set("unrelated:test", {"keep": True}, timeout=60) cache_version_before = get_organization_api_cache_version() result = backfill_all_organization_sources(batch_size=10) self.assertGreaterEqual(result["scanned"], 1) self.assertEqual(result["created_records"], 1) self.assertNotEqual( get_organization_api_cache_version(), cache_version_before, ) self.assertEqual(cache.get("unrelated:test"), {"keep": True}) extension = IndustrialProductionExtension.objects.get( organization=organization, ) record = OrganizationSourceRecord.objects.get(extension=extension) self.assertEqual( record.payload["certificate_number"], "FULL-SOURCE-CERT", ) def test_backfill_parser_batch_task_limits_source_and_batch(self): organization = self._directory_organization( name='ООО "Пакет источника"', inn="7800000402", ogrn="1027700144402", ) IndustrialCertificateRecordFactory( inn=organization.inn, ogrn=organization.ogrn, certificate_number="BATCH-SOURCE-CERT-1", load_batch=1, ) IndustrialCertificateRecordFactory( inn=organization.inn, ogrn=organization.ogrn, certificate_number="BATCH-SOURCE-CERT-2", load_batch=2, ) result = backfill_organization_sources_for_parser_batch( source=ParserLoadLog.Source.INDUSTRIAL, batch_id=2, ) self.assertEqual(result["scanned"], 1) self.assertEqual(result["created_records"], 1) record = OrganizationSourceRecord.objects.get() self.assertEqual(record.payload["certificate_number"], "BATCH-SOURCE-CERT-2") class OrganizationSnapshotScheduleMigrationTest(TestCase): """Checks legacy data migration that schedules the compatibility task.""" def test_migration_seeds_daily_snapshot_refresh_periodic_task(self): migration = import_module( "organizations.migrations.0004_seed_daily_snapshot_refresh_schedule" ) migration.seed_daily_snapshot_refresh_schedule(django_apps, None) migration.seed_daily_snapshot_refresh_schedule(django_apps, None) task = PeriodicTask.objects.get( name=migration.DAILY_ORGANIZATION_SNAPSHOT_TASK_NAME ) self.assertEqual( task.task, "organizations.tasks.refresh_all_organization_data_snapshots", ) self.assertTrue(task.enabled) self.assertEqual(task.args, "[]") self.assertEqual(task.kwargs, '{"batch_size": 100}') self.assertEqual(task.crontab.minute, "30") self.assertEqual(task.crontab.hour, "4") self.assertEqual(str(task.crontab.timezone), "Europe/Moscow") class SourceRecordExportArtifactsTaskTest(TestCase): """Checks nightly artifact generation and its distributed lock.""" def setUp(self): cache.clear() self.export_directory = TemporaryDirectory() self.settings_override = override_settings( SOURCE_RECORD_EXPORT_DIRECTORY=self.export_directory.name, SOURCE_RECORD_EXPORT_GENERATIONS_TO_KEEP=2, SOURCE_RECORD_EXPORT_LOCK_KEY="test:source-record-exports:lock", SOURCE_RECORD_EXPORT_LOCK_TTL_SECONDS=300, ) self.settings_override.enable() def tearDown(self): self.settings_override.disable() self.export_directory.cleanup() cache.clear() super().tearDown() def test_refresh_task_builds_all_artifacts_and_releases_lock(self): result = refresh_source_record_export_artifacts() self.assertEqual(result["status"], "success") self.assertEqual(result["artifacts_count"], 25) self.assertEqual(result["export_year"], timezone.localdate().year) self.assertIsNone(cache.get(settings.SOURCE_RECORD_EXPORT_LOCK_KEY)) def test_refresh_task_skips_when_another_generation_holds_lock(self): cache.set(settings.SOURCE_RECORD_EXPORT_LOCK_KEY, "busy", timeout=300) result = refresh_source_record_export_artifacts() self.assertEqual(result, {"status": "skipped", "reason": "locked"}) class SourceRecordExportScheduleMigrationTest(TestCase): """Checks the nightly Celery Beat schedule for export artifacts.""" def test_migration_seeds_nightly_source_record_export_task(self): migration = import_module( "organizations.migrations.0008_seed_nightly_source_record_exports" ) migration.seed_nightly_source_record_export_schedule(django_apps, None) migration.seed_nightly_source_record_export_schedule(django_apps, None) task = PeriodicTask.objects.get(name=migration.NIGHTLY_SOURCE_EXPORT_TASK_NAME) self.assertEqual( task.task, "organizations.tasks.refresh_source_record_export_artifacts", ) self.assertTrue(task.enabled) self.assertEqual(task.args, "[]") self.assertEqual(task.kwargs, "{}") self.assertEqual(task.crontab.minute, "30") self.assertEqual(task.crontab.hour, "5") self.assertEqual(str(task.crontab.timezone), "Europe/Moscow")