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

from __future__ import annotations

from uuid import uuid4

from acquisition.repositories.postgres import PostgresAcquisitionRepository


class RecordingPostgresRepo(PostgresAcquisitionRepository):
    def __init__(self):
        super().__init__(conn=None)
        self.one_calls = []
        self.one_or_none_calls = []
        self.all_calls = []
        self.one_results = []
        self.one_or_none_results = []
        self.all_results = []

    def _one(self, sql, params):
        self.one_calls.append((sql, params))
        return self.one_results.pop(0)

    def _one_or_none(self, sql, params):
        self.one_or_none_calls.append((sql, params))
        return self.one_or_none_results.pop(0)

    def _all(self, sql, params):
        self.all_calls.append((sql, params))
        return self.all_results.pop(0)


def test_postgres_summary_uses_formal_run_job_item_tables():
    run_id = uuid4()
    batch_id = uuid4()
    query_id = uuid4()
    repo = RecordingPostgresRepo()
    repo.one_results.append(
        {
            "id": run_id,
            "run_key": "r1",
            "batch_id": batch_id,
            "status": "done",
            "query_count": 1,
            "job_count": 3,
            "candidate_count": 7,
            "creation_hit_count": 2,
            "started_at": None,
            "finished_at": None,
        }
    )
    repo.all_results.append(
        [
            {
                "query_id": query_id,
                "query_text": "脚本 开头",
                "job_count": 3,
                "candidate_count": 7,
                "creation_hit_count": 2,
                "platforms": {"xiaohongshu": {"status": "done"}},
            }
        ]
    )

    summary = repo.get_run_summary(run_id)

    first_sql, first_params = repo.one_calls[0]
    second_sql, second_params = repo.all_calls[0]
    assert "FROM acquisition_runs ar" in first_sql
    assert "LEFT JOIN acquisition_jobs aj" in first_sql
    assert "LEFT JOIN candidate_items ci" in first_sql
    assert "LEFT JOIN item_classifications ic" in first_sql
    assert first_params == (run_id,)
    assert "jsonb_object_agg" in second_sql
    assert "WHERE aj.run_id = %s" in second_sql
    assert second_params == (run_id,)
    assert summary["queries"][0]["query_id"] == query_id


def test_postgres_repository_lists_only_creation_candidates_for_decode():
    run_id = uuid4()
    item_id = uuid4()
    repo = RecordingPostgresRepo()
    repo.all_results.append(
        [
            {
                "id": item_id,
                "platform": "xiaohongshu",
                "platform_item_id": "x1",
                "unique_key": "xhs:x1",
                "canonical_url": "https://xhs.test/1",
                "status": "candidate",
            }
        ]
    )

    items = repo.list_creation_candidate_items(run_id=run_id, limit=20)

    sql, params = repo.all_calls[0]
    assert "JOIN acquisition_jobs aj ON aj.id = ci.job_id" in sql
    assert "JOIN item_classifications ic ON ic.item_id = ci.id" in sql
    assert "ic.is_creation_knowledge IS TRUE" in sql
    assert "NOT EXISTS" in sql
    assert "FROM decode_results dr" in sql
    assert "LIMIT %s" in sql
    assert params == (run_id, 20)
    assert items[0].id == item_id
    assert items[0].platform == "xiaohongshu"


def test_postgres_ensure_acquisition_job_merges_metadata_on_conflict():
    job_id = uuid4()
    run_id = uuid4()
    query_id = uuid4()
    repo = RecordingPostgresRepo()
    repo.one_results.append(
        {
            "id": job_id,
            "run_id": run_id,
            "query_id": query_id,
            "platform": "douyin",
            "status": "done",
            "search_limit": 10,
            "display_limit": 5,
            "attempt_count": 1,
            "metadata": {"pagination": {"pages": [{"page_index": 1}]}, "query_text": "q"},
            "error_message": None,
            "started_at": None,
            "finished_at": None,
            "created_at": None,
            "updated_at": None,
        }
    )

    job = repo.ensure_acquisition_job(
        run_id=run_id,
        query_id=query_id,
        platform="douyin",
        search_limit=10,
        display_limit=5,
        status="pending",
        metadata={"query_text": "q"},
    )

    sql, params = repo.one_calls[0]
    assert "metadata = acquisition_jobs.metadata || EXCLUDED.metadata" in sql
    assert params[0] == run_id
    assert params[1] == query_id
    assert job.metadata["pagination"]["pages"][0]["page_index"] == 1


def test_postgres_upsert_candidate_item_prefers_unique_key_for_existing_rows():
    item_id = uuid4()
    repo = RecordingPostgresRepo()
    repo.one_or_none_results.append({"id": item_id})
    repo.one_results.append(
        {
            "id": item_id,
            "job_id": None,
            "query_id": None,
            "platform": "xiaohongshu",
            "platform_item_id": "x1",
            "unique_key": "xhs:x1",
            "canonical_url": "https://xhs.test/1",
            "content_type": "图文",
            "content_mode": "image_post",
            "title": "更新标题",
            "author_name": None,
            "published_at": None,
            "body_text": "完整正文",
            "raw_summary": None,
            "status": "candidate",
            "source_payload": {},
            "metadata": {},
            "error_message": None,
            "created_at": None,
            "updated_at": None,
        }
    )

    item = repo.upsert_candidate_item(
        platform="xiaohongshu",
        platform_item_id="x1",
        unique_key="xhs:x1",
        canonical_url="https://xhs.test/1",
        content_type="图文",
        content_mode="image_post",
        title="更新标题",
        body_text="完整正文",
    )

    lookup_sql, lookup_params = repo.one_or_none_calls[0]
    update_sql, update_params = repo.one_calls[0]
    assert "WHERE unique_key = %s" in lookup_sql
    assert lookup_params == ("xhs:x1",)
    assert "source_payload =" not in update_sql
    assert "unique_key = COALESCE(%s, unique_key)" in update_sql
    assert "content_mode = %s" in update_sql
    assert "body_text = %s" in update_sql
    assert "INSERT INTO candidate_items" not in update_sql
    assert update_params[2] == "xhs:x1"
    assert update_params[5] == "image_post"
    assert update_params[8] == "完整正文"
    assert item.unique_key == "xhs:x1"
    assert item.content_mode == "image_post"
    assert item.body_text == "完整正文"


def test_postgres_query_detail_loads_media_and_classifications_only_when_items_exist():
    run_id = uuid4()
    query_id = uuid4()
    item_id = uuid4()
    repo = RecordingPostgresRepo()
    repo.one_results.append({"id": query_id, "query_text": "q", "status": "ready"})
    repo.all_results.extend(
        [
            [{"id": uuid4(), "run_id": run_id, "query_id": query_id, "platform": "weixin"}],
            [{"id": item_id, "platform": "weixin", "query_id": query_id, "status": "candidate"}],
            [{"id": uuid4(), "item_id": item_id, "media_type": "image", "status": "done"}],
            [{"id": uuid4(), "item_id": item_id, "status": "classified"}],
            [{"item_id": item_id, "decode_status": "decoded", "payload_count": 2}],
        ]
    )

    detail = repo.get_query_detail(run_id=run_id, query_id=query_id)

    assert len(repo.all_calls) == 5
    assert "SELECT * FROM media_assets" in repo.all_calls[2][0]
    assert "SELECT * FROM item_classifications" in repo.all_calls[3][0]
    assert "FROM decode_results dr" in repo.all_calls[4][0]
    assert detail["items"][0]["id"] == item_id
    assert detail["media_assets"][0]["item_id"] == item_id
    assert detail["decode_summaries"][0]["payload_count"] == 2


def test_postgres_latest_query_result_list_uses_lightweight_projection():
    batch_id = uuid4()
    run_id = uuid4()
    query_id = uuid4()
    item_id = uuid4()
    repo = RecordingPostgresRepo()
    repo.one_results.append(
        {"id": query_id, "batch_id": batch_id, "query_text": "q", "status": "ready"}
    )
    repo.one_or_none_results.append({"id": run_id, "batch_id": batch_id, "status": "running"})
    repo.all_results.extend(
        [
            [{"id": uuid4(), "run_id": run_id, "query_id": query_id, "platform": "weixin"}],
            [
                {
                    "id": item_id,
                    "query_id": query_id,
                    "job_id": uuid4(),
                    "platform": "weixin",
                    "title": "标题",
                    "raw_summary": "摘要",
                    "status": "candidate",
                    "content_mode": "article",
                    "metadata": {},
                }
            ],
            [{"id": uuid4(), "item_id": item_id, "media_type": "image", "status": "done"}],
            [
                {
                    "id": uuid4(),
                    "item_id": item_id,
                    "is_creation_knowledge": True,
                    "label": "creation",
                    "confidence": 0.9,
                    "status": "done",
                    "error_message": None,
                }
            ],
            [{"item_id": item_id, "decode_status": "decoded", "particle_count": 1, "payload_count": 1}],
        ]
    )

    detail = repo.get_latest_query_result_list(query_id)

    item_sql = repo.all_calls[1][0]
    media_sql = repo.all_calls[2][0]
    classification_sql = repo.all_calls[3][0]
    assert "SELECT ci.*" not in item_sql
    assert "source_payload" not in item_sql
    assert "body_text" not in item_sql
    assert "LEFT(ci.raw_summary, 700) AS raw_summary" in item_sql
    assert "SELECT DISTINCT ON (item_id)" in media_sql
    assert "media_type IN ('cover', 'image', 'frame')" in media_sql
    assert "SELECT DISTINCT ON (item_id)" in classification_sql
    assert "reason" not in classification_sql
    assert detail["run"]["id"] == run_id
    assert detail["items"][0]["raw_summary"] == "摘要"
    assert detail["media_assets"][0]["item_id"] == item_id
    assert detail["decode_summaries"][0]["particle_count"] == 1


def test_postgres_attach_existing_candidate_updates_query_job_and_metadata():
    item_id = uuid4()
    job_id = uuid4()
    query_id = uuid4()
    repo = RecordingPostgresRepo()
    repo.one_results.append(
        {
            "id": item_id,
            "job_id": job_id,
            "query_id": query_id,
            "platform": "weixin",
            "platform_item_id": "wx1",
            "unique_key": "wx:key",
            "canonical_url": "https://mp.weixin.qq.com/s?...",
            "content_mode": "article",
            "body_text": None,
            "status": "candidate",
            "source_payload": {},
            "metadata": {"acquisition_match_status": "existing"},
            "created_at": None,
            "updated_at": None,
        }
    )

    item = repo.attach_existing_candidate_item(
        item_id,
        job_id=job_id,
        query_id=query_id,
        metadata={"acquisition_match_status": "existing"},
    )

    sql, params = repo.one_calls[0]
    assert "UPDATE candidate_items SET" in sql
    assert "metadata = metadata || %s" in sql
    assert params[0] == job_id
    assert params[1] == query_id
    assert params[3] == item_id
    assert item.metadata["acquisition_match_status"] == "existing"