mostovik-backend/tests/apps/organizations/test_tasks.py

"""Tests for organization source backfill tasks and schedules."""

from importlib import import_module
from tempfile import TemporaryDirectory
from unittest.mock import patch

from apps.parsers.models import ParserLoadLog
from celery.exceptions import Retry
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"], 43)
        self.assertEqual(result["export_year"], timezone.localdate().year)
        self.assertIsNone(cache.get(settings.SOURCE_RECORD_EXPORT_LOCK_KEY))

    def test_refresh_task_retries_when_another_generation_holds_lock(self):
        cache.set(settings.SOURCE_RECORD_EXPORT_LOCK_KEY, "busy", timeout=300)

        with patch.object(
            refresh_source_record_export_artifacts,
            "retry",
            side_effect=Retry("lock held"),
        ) as retry, self.assertRaises(Retry):
            refresh_source_record_export_artifacts()
        retry.assert_called_once_with(countdown=30, max_retries=12)
        self.assertEqual(cache.get(settings.SOURCE_RECORD_EXPORT_LOCK_KEY), "busy")


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")