from __future__ import annotations from uuid import uuid4 import pytest from acquisition.domain import Query, QueryBatch from acquisition.queries.builder import ( QueryBuildOptions, build_creation_query_batch, persist_query_batch, ) from core.config import PgConfig, Settings def _settings() -> Settings: return Settings( pg=PgConfig(host="h", port=5432, user="u", password="p", database="d"), aiddit_crawler_base_url="http://crawler.test", crawler_timeout=30, openrouter_timeout_seconds=90, openrouter_model="m", openrouter_base_url="http://openrouter.test", openrouter_api_key="k", llm_model="m", max_cards=12, frames_dir="f", douyin_ratio="540p", data_dir="", ) class FakeRepo: def __init__(self) -> None: self.batch_kwargs = None self.queries = [] def create_query_batch(self, **kwargs): self.batch_kwargs = kwargs return QueryBatch(id=uuid4(), **kwargs) def add_query(self, **kwargs): self.queries.append(kwargs) return Query(id=uuid4(), **kwargs) def test_build_creation_query_batch_defaults_to_first_two_families(): generated = build_creation_query_batch( _settings(), options=QueryBuildOptions(per=2, batch_n=4, dry=True), ) assert [family["key"] for family in generated["families"]] == ["f1", "f2"] assert generated["metadata"]["active_family_keys"] == ["f1", "f2"] assert sum(len(family["items"]) for family in generated["families"]) == 4 def test_build_creation_query_batch_can_explicitly_enable_reserved_families(): all_keys = ( "f1", "f2", "f4", "f3", "f5", "a_shi", "a_xing", "a_both", "a_purpose", "a_tail", "b_shi", "b_xing", "b_both", "b_purpose", "b_tail", ) generated = build_creation_query_batch( _settings(), options=QueryBuildOptions( per=1, batch_n=4, dry=True, active_family_keys=all_keys, ), ) assert [family["key"] for family in generated["families"]] == list(all_keys) assert generated["metadata"]["active_family_keys"] == list(all_keys) def test_build_creation_query_batch_rejects_unknown_family_key(): with pytest.raises(ValueError, match="unknown query family"): build_creation_query_batch( _settings(), options=QueryBuildOptions( dry=True, active_family_keys=("f1", "no_such_family"), ), ) def test_persist_query_batch_writes_formal_batch_and_query_contract(): repo = FakeRepo() generated = { "metadata": { "query_filter_prompt_version": "abc123", "active_family_keys": ["f1", "f2"], }, "families": [ { "key": "f1", "name": "实质 × 模态 × 业务阶段 × 知识类型", "axes": ["实质", "模态"], "items": [ { "query": "反转 视频 脚本 怎么做", "parts": {"实质": "反转", "模态": "视频"}, "keep": True, "valid": 8, "relevant": True, "reason": "可搜创作方法", } ], } ], } batch, count = persist_query_batch(repo, generated, name="formal-demo") assert batch.name == "formal-demo" assert count == 1 assert repo.batch_kwargs["status"] == "ready" assert repo.batch_kwargs["target_platforms"] == ["xiaohongshu", "weixin", "douyin"] assert repo.batch_kwargs["metadata"]["active_family_keys"] == ["f1", "f2"] row = repo.queries[0] assert row["batch_id"] == batch.id assert row["query_text"] == "反转 视频 脚本 怎么做" assert row["axes"] == {"实质": "反转", "模态": "视频"} assert row["keep"] is True assert row["filter_reason"] == "可搜创作方法" assert row["metadata"]["family_key"] == "f1" assert row["metadata"]["valid"] == 8 assert row["metadata"]["query_filter_prompt_version"] == "abc123"