script-build-and-framework-tmp/script_build_host/tests/test_decode_index_builder.py

from __future__ import annotations

import asyncio
import json
from hashlib import sha256
from pathlib import Path

import httpx
import numpy as np
import pytest

from script_build_host.infrastructure.decode_index_builder import (
    DashScopeEmbeddingClient,
    EmbeddingApiError,
    extract_decode_chunks,
    rebuild_decode_index,
)


def _raw_fixture(root: Path) -> Path:
    raw = root / "script_decode_raw_data" / "账号甲"
    raw.mkdir(parents=True)
    (raw / "post-1.json").write_text(
        json.dumps(
            {
                "帖子ID": "post-1",
                "段落列表": [
                    {
                        "id": "p1",
                        "名称": "开头",
                        "主题原子点": [
                            {
                                "维度名": "对象",
                                "维度名描述": "面向谁",
                                "维度值": "创业者",
                                "维度值描述": "早期团队",
                            }
                        ],
                        "子项": [
                            {
                                "id": "p1.1",
                                "名称": "子段",
                                "形式原子点": [{"维度": "表达", "原子点": "口播"}],
                            }
                        ],
                    }
                ],
            },
            ensure_ascii=False,
        ),
        encoding="utf-8",
    )
    return raw.parent


def _sha256(path: Path) -> str:
    return sha256(path.read_bytes()).hexdigest()


def test_extract_decode_chunks_matches_legacy_fields_and_nested_paragraphs(
    tmp_path: Path,
) -> None:
    raw_root = _raw_fixture(tmp_path)
    chunks = extract_decode_chunks(raw_root)

    assert len(chunks) == 6
    assert [chunk["text_field"] for chunk in chunks] == [
        "维度名",
        "维度名描述",
        "维度值",
        "维度值描述",
        "维度",
        "原子点",
    ]
    assert chunks[-1]["paragraph_id"] == "p1.1"
    assert chunks[-1]["json_relpath"] == "script_decode_raw_data/账号甲/post-1.json"


@pytest.mark.asyncio
async def test_rebuild_decode_index_is_complete_resumable_and_atomic(tmp_path: Path) -> None:
    raw_root = _raw_fixture(tmp_path)
    source_path = raw_root / "账号甲" / "post-1.json"
    source_before = _sha256(source_path)
    index_root = tmp_path / "index"
    index_root.mkdir()
    (index_root / "old-marker").write_text("old", encoding="utf-8")
    request_sizes: list[int] = []

    def handler(request: httpx.Request) -> httpx.Response:
        payload = json.loads(request.content)
        texts = payload["input"]
        request_sizes.append(len(texts))
        data = [
            {"index": index, "embedding": [1.0, float(index + 1), 0.5, 0.25]}
            for index in reversed(range(len(texts)))
        ]
        return httpx.Response(
            200,
            json={"model": "text-embedding-v4", "data": data, "usage": {"prompt_tokens": 7}},
        )

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http:
        client = DashScopeEmbeddingClient(
            http,
            api_key="secret",
            endpoint="https://dashscope.example/embeddings",
            model="text-embedding-v4",
            dimension=4,
            concurrency=2,
        )
        result = await rebuild_decode_index(
            raw_root=raw_root,
            index_root=index_root,
            client=client,
            batch_size=2,
            shard_size=4,
        )

    assert request_sizes == [2, 2, 2]
    assert not (index_root / "old-marker").exists()
    assert _sha256(source_path) == source_before
    assert result.manifest["provider"] == "aliyun-bailian"
    assert result.manifest["chunk_count"] == 6
    assert result.manifest["request_count"] == 3
    assert result.manifest["total_tokens"] == 21
    assert np.load(index_root / "embeddings.npy").shape == (6, 4)
    assert len((index_root / "meta.jsonl").read_text(encoding="utf-8").splitlines()) == 6
    assert not await asyncio.to_thread(lambda: list(tmp_path.glob(".index-build-*")))


@pytest.mark.asyncio
async def test_failed_rebuild_keeps_previously_published_index(tmp_path: Path) -> None:
    raw_root = _raw_fixture(tmp_path)
    index_root = tmp_path / "index"
    index_root.mkdir()
    (index_root / "old-marker").write_text("old", encoding="utf-8")

    def handler(_: httpx.Request) -> httpx.Response:
        return httpx.Response(400, json={"message": "bad request"})

    async with httpx.AsyncClient(transport=httpx.MockTransport(handler)) as http:
        client = DashScopeEmbeddingClient(
            http,
            api_key="secret",
            endpoint="https://dashscope.example/embeddings",
            model="text-embedding-v4",
            dimension=4,
            max_retries=0,
        )
        with pytest.raises(EmbeddingApiError, match="HTTP 400"):
            await rebuild_decode_index(
                raw_root=raw_root,
                index_root=index_root,
                client=client,
                batch_size=2,
                shard_size=4,
            )

    assert (index_root / "old-marker").read_text(encoding="utf-8") == "old"
    assert await asyncio.to_thread(lambda: list(tmp_path.glob(".index-build-*")))