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