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

from __future__ import annotations

from uuid import UUID, uuid4

import pytest

from acquisition.domain import CandidateItem, MediaAsset
from acquisition.runner import RunBatchResult
from core.models import Post
from decode_content.models import GateResult, ReadResult
from decode_content.service import DecodeWorkflowOutput
from pipeline import acquisition_runner, creation_pipeline
from pipeline.acquisition_runner import AcquisitionStageResult
from pipeline.decode_runner import run_decode_stage
from pipeline.decode_runner import DecodeBatchResult
from pipeline.dedupe import dedupe_candidate_items, item_dedupe_key, should_decode_item
from pipeline.models import PipelineJob, PipelineRun


class FakePipelineRepo:
    def __init__(self):
        self.jobs = []

    def create_pipeline_run(self, **kwargs):
        return PipelineRun(id=uuid4(), **kwargs)

    def save_pipeline_job(self, **kwargs):
        job = PipelineJob(id=uuid4(), **kwargs)
        self.jobs.append(job)
        return job

    def mark_job_status(self, job_id, **kwargs):
        job = next(row for row in self.jobs if row.id == job_id)
        updated = job.model_copy(update=kwargs)
        self.jobs.append(updated)
        return updated

    def get_resume_cursor(self, run_id):
        return None


def test_dedupe_candidate_items_prefers_unique_key_then_platform_item_id_then_url():
    first = CandidateItem(
        id=uuid4(),
        platform="xiaohongshu",
        platform_item_id="x1",
        unique_key="xhs:x1",
        canonical_url="https://a",
    )
    dup = CandidateItem(
        id=uuid4(),
        platform="xiaohongshu",
        platform_item_id="x1",
        unique_key="xhs:x1",
        canonical_url="https://b",
    )
    platform_id_only = CandidateItem(id=uuid4(), platform="douyin", platform_item_id="dy1")
    url_only = CandidateItem(id=uuid4(), platform="weixin", canonical_url="HTTPS://MP.TEST/A")

    assert item_dedupe_key(first) == "unique_key:xhs:x1"
    assert item_dedupe_key(platform_id_only) == "douyin:id:dy1"
    assert item_dedupe_key(url_only) == "url:https://mp.test/a"
    assert dedupe_candidate_items([first, dup, platform_id_only, url_only]) == [
        first,
        platform_id_only,
        url_only,
    ]
    assert should_decode_item(first, decoded_item_ids={str(first.id)}).keep is False


def test_acquisition_stage_wraps_existing_runner(monkeypatch):
    batch_id = uuid4()
    pipeline_run = PipelineRun(id=uuid4(), status="running", current_stage="search")
    pipeline_repo = FakePipelineRepo()

    def fake_run_batch(acquisition_repo, *, batch_id, settings, **kwargs):
        return RunBatchResult(run_id=uuid4(), total_jobs=1, done=1, partial=0, failed=0, skipped=0)

    monkeypatch.setattr(acquisition_runner, "run_batch", fake_run_batch)
    result = acquisition_runner.run_acquisition_stage(
        acquisition_repo=object(),
        batch_id=batch_id,
        settings=object(),
        pipeline_repo=pipeline_repo,
        pipeline_run=pipeline_run,
    )

    assert result.acquisition.done == 1
    assert pipeline_repo.jobs[0].stage == "search"
    assert pipeline_repo.jobs[-1].status == "done"


def test_acquisition_stage_marks_pipeline_job_failed_and_reraises(monkeypatch):
    pipeline_run = PipelineRun(id=uuid4(), status="running", current_stage="search")
    pipeline_repo = FakePipelineRepo()

    def fake_run_batch(*args, **kwargs):
        raise RuntimeError("acquisition exploded")

    monkeypatch.setattr(acquisition_runner, "run_batch", fake_run_batch)

    with pytest.raises(RuntimeError, match="acquisition exploded"):
        acquisition_runner.run_acquisition_stage(
            acquisition_repo=object(),
            batch_id=uuid4(),
            settings=object(),
            pipeline_repo=pipeline_repo,
            pipeline_run=pipeline_run,
        )

    assert pipeline_repo.jobs[0].status == "running"
    assert pipeline_repo.jobs[-1].status == "failed"
    assert pipeline_repo.jobs[-1].error_message == "acquisition exploded"


class FakeCandidateRepo:
    def __init__(self):
        self.item_id = uuid4()
        self.items = [
            CandidateItem(
                id=self.item_id,
                platform="xiaohongshu",
                platform_item_id="x1",
                canonical_url="https://xhs/1",
                title="脚本",
                status="candidate",
            ),
            CandidateItem(
                id=uuid4(),
                platform="xiaohongshu",
                platform_item_id="x1",
                title="重复",
                status="candidate",
            ),
        ]

    def list_creation_candidate_items(self, *, run_id: UUID | None = None, limit: int = 100):
        return self.items

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


class FakeDecodeService:
    def __init__(self):
        self.posts: list[Post] = []

    def decode_post(self, *, item_id: UUID, post: Post):
        self.posts.append(post)
        return DecodeWorkflowOutput(
            item_id=item_id,
            read_result=ReadResult(text="ok"),
            gate_result=GateResult(passed=True),
            status="decoded",
        )


def test_decode_stage_loads_candidate_media_and_dedupes():
    repo = FakeCandidateRepo()
    service = FakeDecodeService()

    result = run_decode_stage(candidate_repo=repo, decode_service=service)

    assert result.total == 1
    assert result.decoded == 1
    assert result.failed == 0
    assert service.posts[0].image_urls == ["https://cdn.test/1.jpg"]


def test_creation_pipeline_orchestrates_acquisition_then_decode(monkeypatch):
    batch_id = uuid4()
    acquisition_run_id = uuid4()
    pipeline_repo = FakePipelineRepo()
    calls = {}

    def fake_acquisition_stage(**kwargs):
        calls["acquisition"] = kwargs
        return AcquisitionStageResult(
            pipeline_job=None,
            acquisition=RunBatchResult(
                run_id=acquisition_run_id,
                total_jobs=3,
                done=3,
                partial=0,
                failed=0,
                skipped=0,
            ),
        )

    def fake_decode_stage(**kwargs):
        calls["decode"] = kwargs
        return DecodeBatchResult(total=2, decoded=2, skipped=0, failed=0, outputs=[])

    monkeypatch.setattr(creation_pipeline, "run_acquisition_stage", fake_acquisition_stage)
    monkeypatch.setattr(creation_pipeline, "run_decode_stage", fake_decode_stage)

    result = creation_pipeline.run_creation_pipeline(
        acquisition_repo=object(),
        batch_id=batch_id,
        settings=object(),
        pipeline_repo=pipeline_repo,
        decode_service=object(),
        decode_limit=20,
    )

    assert result.pipeline_run is not None
    assert result.pipeline_run.batch_id == batch_id
    assert result.acquisition.acquisition.done == 3
    assert result.decode is not None
    assert result.decode.decoded == 2
    assert calls["acquisition"]["pipeline_run"] == result.pipeline_run
    assert calls["decode"]["run_id"] == acquisition_run_id
    assert calls["decode"]["limit"] == 20