Create-knowledge-find-decode/tests/test_pipeline_tracing.py

from __future__ import annotations

from pathlib import Path
from contextlib import nullcontext
from uuid import uuid4

from acquisition.domain import CandidateItem, MediaAsset
from core.models import Post
from decode_content.models import GateResult, ReadResult
from decode_content.service import DecodeWorkflowOutput
from pipeline.decode_runner import run_decode_stage
from pipeline.tracing import TraceConfig, TraceContext, sanitize_payload
from scripts import backfill_pipeline_traces


def test_formal_schema_declares_append_only_trace_tables():
    text = Path("db/migrations/001_creation_knowledge_schema.sql").read_text(encoding="utf-8")

    for table in (
        "pipeline_runs",
        "pipeline_jobs",
        "pipeline_run_events",
        "candidate_item_hits",
        "llm_call_traces",
    ):
        assert f"CREATE TABLE IF NOT EXISTS creation_knowledge.{table}" in text
    assert "ALTER TABLE creation_knowledge.candidate_items" not in text
    assert not Path("db/migrations/005_pipeline_trace_schema.sql").exists()


def test_formal_schema_uses_deduped_contract_artifacts():
    text = Path("db/migrations/001_creation_knowledge_schema.sql").read_text(encoding="utf-8")

    assert "CREATE TABLE IF NOT EXISTS creation_knowledge.contract_artifacts" in text
    assert "CREATE TABLE IF NOT EXISTS creation_knowledge.run_contract_artifacts" in text
    assert "UNIQUE (contract_name, contract_type, source_path, content_hash)" in text
    assert "CREATE TABLE IF NOT EXISTS creation_knowledge.contract_snapshots" not in text
    assert not Path("db/migrations/005_pipeline_trace_schema.sql").exists()


def test_postgres_decode_repository_writes_deduped_contract_artifacts():
    text = Path("decode_content/repositories/postgres.py").read_text(encoding="utf-8")

    assert "INSERT INTO contract_artifacts" in text
    assert "INSERT INTO run_contract_artifacts" in text
    assert "ON CONFLICT (contract_name, contract_type, source_path, content_hash)" in text
    assert "INSERT INTO contract_snapshots" not in text


def test_sanitize_payload_omits_base64_data_urls_but_keeps_metadata():
    data_url = "data:video/mp4;base64," + ("a" * 120)

    payload = sanitize_payload(
        {"video": data_url, "url": "https://cdn.test/video.mp4"},
        max_chars=50,
    )

    assert payload["url"] == "https://cdn.test/video.mp4"
    assert payload["video"]["omitted"] == "base64_data_url"
    assert payload["video"]["chars"] == len(data_url)
    assert payload["video"]["sha256"]


def test_trace_config_invalid_integer_falls_back_to_default(monkeypatch):
    monkeypatch.setenv("CK_TRACE_MAX_PAYLOAD_CHARS", "not-a-number")

    config = TraceConfig.from_env(".env.missing")

    assert config.max_payload_chars == 50_000


class CollectingTraceWriter:
    def __init__(self):
        self.events = []
        self.hits = []
        self.calls = []

    def event(self, *, context, stage, event_type, **kwargs):
        self.events.append({"context": context, "stage": stage, "event_type": event_type, **kwargs})

    def candidate_hit(self, *, context, **kwargs):
        self.hits.append({"context": context, **kwargs})

    def llm_call(self, *, context, stage, substage=None, **kwargs):
        self.calls.append({"context": context, "stage": stage, "substage": substage, **kwargs})


class CandidateRepo:
    def __init__(self):
        self.item_id = uuid4()
        self.job_id = uuid4()
        self.query_id = uuid4()

    def list_creation_candidate_items(self, *, run_id=None, limit=100):
        return [
            CandidateItem(
                id=self.item_id,
                job_id=self.job_id,
                query_id=self.query_id,
                platform="xiaohongshu",
                platform_item_id="x1",
                unique_key="xhs:x1",
                canonical_url="https://xhs.test/1",
                title="拍照姿势",
                status="candidate",
            )
        ]

    def list_media_assets_for_item(self, item_id):
        return [
            MediaAsset(
                id=uuid4(),
                item_id=item_id,
                media_type="image",
                cdn_url="https://cdn.test/1.jpg",
                position=1,
                status="done",
            )
        ]


class DecodeService:
    def decode_post(self, *, item_id, post: Post):
        assert post.image_urls == ["https://cdn.test/1.jpg"]
        return DecodeWorkflowOutput(
            item_id=item_id,
            read_result=ReadResult(text="ok"),
            gate_result=GateResult(passed=True),
            status="decoded",
        )


def test_decode_stage_emits_trace_events_without_changing_result():
    writer = CollectingTraceWriter()
    repo = CandidateRepo()
    pipeline_run_id = uuid4()

    result = run_decode_stage(
        candidate_repo=repo,
        decode_service=DecodeService(),
        trace_writer=writer,
        trace_context=TraceContext(pipeline_run_id=pipeline_run_id, stage="decode"),
    )

    assert result.decoded == 1
    event_types = [event["event_type"] for event in writer.events]
    assert event_types[0] == "decode_stage_started"
    assert "decode_item_started" in event_types
    assert "decode_item_finished" in event_types
    assert event_types[-1] == "decode_stage_finished"
    assert writer.events[1]["context"].pipeline_run_id == pipeline_run_id


def test_backfill_pipeline_traces_dry_run_does_not_write(monkeypatch):
    acquisition_run_id = uuid4()

    monkeypatch.setattr(backfill_pipeline_traces.CreationDbConfig, "from_env", lambda _env_file: object())
    monkeypatch.setattr(backfill_pipeline_traces, "transaction", lambda _config: nullcontext(object()))
    monkeypatch.setattr(
        backfill_pipeline_traces,
        "_acquisition_runs",
        lambda _conn, limit=None: [
            {
                "id": acquisition_run_id,
                "run_key": "manual-api:test",
                "status": "done",
                "job_count": 3,
                "candidate_count": 12,
            }
        ],
    )

    result = backfill_pipeline_traces.backfill(env_file=".env", apply=False)

    assert result["dry_run"] is True
    assert result["acquisition_run_count"] == 1
    assert result["candidate_hit_count"] == 12
    assert result["runs"][0]["acquisition_run_id"] == str(acquisition_run_id)