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

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"