zhifu-platform/tests/bounded_capacity_acceptance.py

import asyncio
import json
import os
import platform
from pathlib import Path

import httpx
from capacity_acceptance import accounts, measure
from retrieval_evaluation import publish_dataset
from sqlalchemy import func, select
from test_after_sales import scenario  # noqa: F401
from test_knowledge_memory import cleanup_vectors

from zhifu.modules.agent_runtime.models import ModelAttempt


async def test_bounded_current_api_and_retrieval_capacity(scenario):  # noqa: F811
    report = {
        "environment": platform.platform(), "python": platform.python_version(),
        "logical_cpus": os.cpu_count(), "api_workers": 1, "database_pool": "5+5",
        "inference_device": "cpu", "documents": 12, "account_count": 16,
        "maximum_load_requests": 3200, "maximum_seconds": 120, "stages": [],
        "production_acceptance": "NOT_RUN",
        "production_targets": {"api_rps": 1000, "api_p95_ms": 300,
                               "error_rate_below": 0.001, "sse_connections": 3000,
                               "agent_turns_per_second": 50, "first_response_p95_ms": 3000},
    }
    try:
        async with asyncio.timeout(120):
            await publish_dataset(scenario)
            identities = await accounts(scenario)
            async with httpx.AsyncClient(
                base_url="http://127.0.0.1:18000", timeout=15, trust_env=False,
                limits=httpx.Limits(max_connections=32),
            ) as client:
                for workload, method, path, body in (
                    ("orders", "GET", f"/openapi/v1/orders?store_id={scenario['store']}", None),
                    ("knowledge", "POST",
                     f"/openapi/v1/stores/{scenario['store']}/knowledge/search",
                     {"query": "键盘脏了能放水里洗吗"}),
                ):
                    for concurrency in (1, 16):
                        result = await measure(client, method, path, body, concurrency,
                                               identities, seconds=5, maximum_requests=800)
                        report["stages"].append({"workload": workload, **result})
                        assert set(result["statuses"]) <= {"200", "429"}, result
                        assert result["statuses"].get("200", 0) > 0
                        assert result["requests"] <= 800
            async with scenario["app"].state.sessions() as session:
                report["generation_attempts"] = await session.scalar(
                    select(func.count()).select_from(ModelAttempt).where(
                        ModelAttempt.tenant_id == scenario["tenant"].bytes))
            assert report["generation_attempts"] == 0
    finally:
        await asyncio.to_thread(Path(".work/long-capacity.json").write_text,
            json.dumps(report, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
        await cleanup_vectors(scenario["tenant"])