182 lines
6.8 KiB
Python
182 lines
6.8 KiB
Python
"""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")
|