DataOps-platform/tests/data_research/test_catalog_execution.py

from __future__ import annotations

from dataclasses import asdict

import pytest

from tests.data_research.test_ingestion_service import (
    MemoryJobRepository,
    payload,
)


class MemorySnapshotRepository:
    def __init__(self):
        self.records = {}

    def find(self, job_uid, attempt):
        return self.records.get((str(job_uid), int(attempt)))

    def persist(self, job_uid, attempt, snapshot):
        from app.core.data_research.catalog.models import CatalogSnapshotRecord

        key = (str(job_uid), int(attempt))
        existing = self.records.get(key)
        if existing is not None:
            return existing
        field_count = sum(len(asset.fields) for asset in snapshot.assets)
        record = CatalogSnapshotRecord(
            uid=f"snapshot-{attempt}",
            job_uid=str(job_uid),
            source_uid=snapshot.data_source_uid,
            attempt=int(attempt),
            database_type=snapshot.database_type,
            content_hash=snapshot.content_hash,
            snapshot=asdict(snapshot),
            evidence_count=field_count,
        )
        self.records[key] = record
        return record

    def list(self, job_uid):
        return [
            record
            for (candidate_uid, _attempt), record in sorted(self.records.items())
            if candidate_uid == str(job_uid)
        ]


def catalog_snapshot():
    from app.core.data_research.catalog.models import (
        CatalogAsset,
        CatalogField,
        CatalogSnapshot,
    )

    return CatalogSnapshot(
        data_source_uid="00000000-0000-0000-0000-000000000001",
        database_type="postgresql",
        assets=(
            CatalogAsset(
                key="source:asset.equipment",
                schema="asset",
                name="equipment",
                asset_type="table",
                fields=(
                    CatalogField(
                        key="source:asset.equipment.equipment_code",
                        schema="asset",
                        asset="equipment",
                        name="equipment_code",
                        ordinal_position=1,
                        data_type="varchar",
                        nullable=False,
                    ),
                    CatalogField(
                        key="source:asset.equipment.location_code",
                        schema="asset",
                        asset="equipment",
                        name="location_code",
                        ordinal_position=2,
                        data_type="varchar",
                        nullable=True,
                    ),
                ),
            ),
        ),
    )


class CatalogCollector:
    def __init__(self, result=None, error=None):
        self.result = result or catalog_snapshot()
        self.error = error
        self.calls = []

    def collect(self, source_uid, scope):
        self.calls.append((source_uid, scope))
        if self.error is not None:
            raise self.error
        return self.result


def catalog_payload(**overrides):
    return payload(
        artifact_uid=None,
        job_type="catalog_collect",
        parser_version="catalog-v1",
        parameters={
            "include_schemas": ["asset"],
            "exclude_schemas": ["pg_catalog"],
            "include_tables": ["equipment"],
            "exclude_tables": [],
        },
        **overrides,
    )


def test_catalog_execution_persists_snapshot_evidence_and_job_report():
    from app.core.data_research.catalog.execution import CatalogIngestionExecutor
    from app.core.data_research.ingestion import IngestionService

    jobs = MemoryJobRepository()
    ingestion = IngestionService(jobs)
    job, _ = ingestion.create_job(catalog_payload(), actor_uid="editor-1")
    collector = CatalogCollector()
    snapshots = MemorySnapshotRepository()
    executor = CatalogIngestionExecutor(ingestion, collector, snapshots)

    completed = executor.execute(job.uid)

    assert completed.status == "awaiting_review"
    assert completed.attempt_count == 1
    assert completed.statistics == {
        "snapshot_uid": "snapshot-1",
        "content_hash": catalog_snapshot().content_hash,
        "asset_count": 1,
        "field_count": 2,
        "evidence_count": 2,
    }
    assert snapshots.find(job.uid, 1).attempt == 1
    assert collector.calls[0][1].include_schemas == ("asset",)

    replayed = executor.execute(job.uid)
    assert replayed == completed
    assert len(collector.calls) == 1


def test_failed_catalog_collection_records_stage_and_can_retry():
    from app.core.data_research.catalog.execution import CatalogIngestionExecutor
    from app.core.data_research.ingestion import IngestionService

    jobs = MemoryJobRepository()
    ingestion = IngestionService(jobs)
    job, _ = ingestion.create_job(catalog_payload(), actor_uid="editor-1")
    collector = CatalogCollector(
        error=RuntimeError("password=clear-secret connection refused")
    )
    snapshots = MemorySnapshotRepository()
    executor = CatalogIngestionExecutor(ingestion, collector, snapshots)

    with pytest.raises(RuntimeError, match="connection refused"):
        executor.execute(job.uid)

    failed = ingestion.get_job(job.uid)
    assert failed.status == "failed"
    assert failed.attempt_count == 1
    assert failed.failure_stage == "extracting"
    assert "clear-secret" not in failed.last_error
    assert snapshots.list(job.uid) == []

    ingestion.retry(job.uid)
    collector.error = None
    completed = executor.execute(job.uid)

    assert completed.status == "awaiting_review"
    assert completed.attempt_count == 2
    assert snapshots.find(job.uid, 2).attempt == 2


def test_catalog_scope_rejects_non_list_or_non_text_values_before_collection():
    from app.core.data_research.catalog.execution import catalog_scope_from_parameters
    from app.core.data_research.errors import IngestionPayloadInvalid

    with pytest.raises(IngestionPayloadInvalid, match="include_schemas"):
        catalog_scope_from_parameters({"include_schemas": "asset"})
    with pytest.raises(IngestionPayloadInvalid, match="include_tables"):
        catalog_scope_from_parameters({"include_tables": ["equipment", 42]})