vent_agent_2.0/tests/test_shaanxi_gateway.py

"""取数网关协议、多业务身份、文件下载及旧SDK协议兼容性。"""
import asyncio
from dataclasses import replace
import hashlib
import json

import httpx
import pytest

from db import business_store as store, db_conn
from extensions import registry as extensions
from extensions.registry import ExtensionRegistry, GatewayBinding
from extensions.main_fan.schemas import MainFanInput
from extensions.shaanxi.client import GatewayClient, GatewayError, GatewayStopped
from extensions.shaanxi.settings import GatewaySettings
from runtime.business_models import JudgmentResult, RemoteFile
from runtime.business_tasks import BusinessTaskController
from runtime.engine import ExecutionResult


def settings(**kwargs):
    return GatewaySettings("https://gateway.invalid", "gateway-secret", "https://files.invalid", **kwargs)


def response(data=None, **kwargs):
    return httpx.Response(200, json={"code": 200, "success": True, "data": data}, **kwargs)


@pytest.mark.asyncio
async def test_concurrent_calls_keep_agent_identity_and_registration_contract():
    seen = []
    async def handler(request):
        seen.append((request.url.path, request.headers["X-Agent-Code"], json.loads(request.content) if request.content else None))
        assert request.headers["Authorization"] == "Bearer gateway-secret"
        await asyncio.sleep(0)
        return response([])
    client = GatewayClient(settings(), transport=httpx.MockTransport(handler))
    try:
        await asyncio.gather(client.query("agent-one", "monitor_platform", "alarms"),
                             client.query("agent-two", "monitor_platform", "alarm_detail"))
        assert {item[1] for item in seen} == {"agent-one", "agent-two"}
        original = extensions.load_extensions().business("main_fan")
        profile = replace(original, gateway=GatewayBinding("fan-code", "fan-check", topics=()))
        await client.register(profile)
        assert seen[-1][2]["agentCode"] == "fan-code"
        assert seen[-1][2]["topics"] == []
        assert seen[-1][2]["enabled"] is True
    finally:
        await client.close()


@pytest.mark.asyncio
async def test_http_200_business_failure_and_nested_proxy_status_are_rejected():
    answers = [httpx.Response(200, json={"code": 500, "success": False}),
               response({"status_code": 500, "body": {}}),
               response({"status_code": 200, "body": {"code": 500}}),
               httpx.Response(500, json={"detail": "智能体已停止"})]
    client = GatewayClient(settings(), transport=httpx.MockTransport(lambda r: answers.pop(0)))
    try:
        with pytest.raises(GatewayError):
            await client.push("one", {"taskId": "x"})
        with pytest.raises(GatewayError) as error:
            await client.proxy("one", "/query")
        assert error.value.retryable
        with pytest.raises(GatewayError):
            await client.proxy("one", "/query")
        with pytest.raises(GatewayStopped):
            await client.token("one")
    finally:
        await client.close()


@pytest.mark.asyncio
async def test_download_uses_platform_headers_and_rotates_expired_token(tmp_path):
    tokens = ["expired", "fresh"]
    downloads = []
    def handler(request):
        if request.url.host == "gateway.invalid":
            assert request.headers["X-Agent-Code"] == "fan-code"
            return response({"access_token": tokens.pop(0), "client_id": "platform-client"})
        downloads.append(request.headers["Authorization"])
        assert request.headers["clientid"] == "platform-client"
        assert "X-Agent-Code" not in request.headers
        if downloads[-1] == "Bearer expired":
            return httpx.Response(401)
        return httpx.Response(200, content=b"%PDF-test", headers={"content-type": "application/pdf"})
    client = GatewayClient(settings(), transport=httpx.MockTransport(handler))
    try:
        source = RemoteFile(file_id="r1", version="v1", filename="report.pdf", path="/download", params={"id": "r1"})
        metadata = await client.download("fan-code", source, tmp_path / "report.pdf")
        assert downloads == ["Bearer expired", "Bearer fresh"]
        assert metadata["sha256"] == hashlib.sha256(b"%PDF-test").hexdigest()
        assert metadata["version"] == "v1"
        assert not (tmp_path / "report.pdf.part").exists()
    finally:
        await client.close()


@pytest.mark.asyncio
@pytest.mark.parametrize("failure", ["same_token", "too_large", "json"])
async def test_bad_download_does_not_leave_a_file(tmp_path, failure):
    def handler(request):
        if request.url.host == "gateway.invalid":
            return response({"access_token": "old", "client_id": "c"})
        if failure == "same_token":
            return httpx.Response(401)
        if failure == "json":
            return httpx.Response(200, json={"code": 401})
        return httpx.Response(200, content=b"oversized")
    client = GatewayClient(settings(max_file_bytes=4), transport=httpx.MockTransport(handler))
    target = tmp_path / "report.pdf"
    try:
        with pytest.raises(GatewayError):
            await client.download("one", RemoteFile(file_id="r", version="v", filename="r.pdf", path="/file"), target)
        assert not target.exists()
        assert not target.with_suffix(".pdf.part").exists()
    finally:
        await client.close()


@pytest.mark.asyncio
async def test_sse_frames_and_stop_signal():
    frames = 'event: fan.task\nid: event-1\ndata: {"data":{"mineCode":"m"}}\n\n: keepalive\n\n'
    responses = [httpx.Response(200, text=frames, headers={"content-type": "text/event-stream"}), httpx.Response(403)]
    client = GatewayClient(settings(), transport=httpx.MockTransport(lambda r: responses.pop(0)))
    try:
        events = [item async for item in client.events("fan-code")]
        assert events[0] is None
        assert events[1].message_id == "event-1" and events[1].topic == "fan.task"
        with pytest.raises(GatewayStopped):
            async for _ in client.events("fan-code"):
                pass
    finally:
        await client.close()


@pytest.mark.asyncio
async def test_rate_limit_hint_is_preserved_and_incomplete_receipt_rejected():
    replies = [httpx.Response(429, headers={"Retry-After": "120"}), response({})]
    client = GatewayClient(settings(), transport=httpx.MockTransport(lambda r: replies.pop(0)))
    try:
        with pytest.raises(GatewayError) as limited:
            await client.push("fan-code", {"taskId": "one"})
        assert limited.value.retryable and limited.value.retry_after == 120
        with pytest.raises(GatewayError, match="回执"):
            await client.push("fan-code", {"taskId": "one"})
    finally:
        await client.close()


@pytest.mark.asyncio
async def test_downloaded_material_is_visible_to_business_in_its_own_workspace(monkeypatch, tmp_path):
    from agents import workspace
    from extensions.main_fan.tools import inspect_main_fan_materials
    base = extensions.load_extensions()
    original = base.business_agents["main_fan"]
    monkeypatch.setitem(base.business_agents, "main_fan", replace(original, gateway=GatewayBinding("fan-code", "fan-check")))
    monkeypatch.setattr(workspace.config, "ROOT", tmp_path)
    def handler(request):
        if request.url.host == "gateway.invalid":
            return response({"access_token": "file-token", "client_id": "client-id"})
        return httpx.Response(200, content=b"report", headers={"content-type": "application/pdf"})
    received = []
    async def executor(**kwargs):
        received.append(kwargs)
        value = inspect_main_fan_materials([{"path": "uploaded"}], "m", "矿")
        return ExecutionResult("run", kwargs["session_id"], "success", "", "", {}, [], [], value)
    client = GatewayClient(settings(), transport=httpx.MockTransport(handler))
    try:
        source = {"file_id": "f", "version": "1", "filename": "report.pdf", "path": "/download"}
        job, _ = store.enqueue("main_fan", "file-owner", {"mine_code": "m", "mine_name": "矿", "materials": []}, [source])
        await BusinessTaskController(client, executor=executor).run_job(store.claim("main_fan"))
        saved = store.get(job["id"])
        assert saved["status"] == "needs_materials"
        assert "/workspace/uploads/" in received[0]["message"]
        metadata = db_conn.query_one("SELECT * FROM business_files WHERE job_id = ?", (job["id"],))
        assert metadata["sha256"] == hashlib.sha256(b"report").hexdigest()
        assert metadata["local_path"].startswith(str(workspace.session_root("file-owner", saved["session_id"])))
    finally:
        await client.close()