tests/test_data_routing_and_export_guard.py

import ast
import re
import unittest
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]


def load_helpers():
    source = ROOT / "ui_server.py"
    tree = ast.parse(source.read_text(encoding="utf-8"))
    names = {
        "_is_fisheries_prompt",
        "_needs_ocean_mcp",
        "_is_ocean_export_request",
        "_apply_ocean_export_defaults",
        "_ocean_export_execution_error",
    }
    functions = [
        node for node in tree.body
        if isinstance(node, ast.FunctionDef) and node.name in names
    ]
    catalog = {}
    exec((ROOT / "sidebar_catalog.py").read_text(encoding="utf-8"), catalog)
    namespace = {
        "re": re,
        "_FISHERY_TERMS": catalog["FISHERY_TERMS"],
        "_OCEAN_TOOL_TERMS": catalog["OCEAN_TOOL_TERMS"],
        "_OCEAN_STATUS_TERMS": catalog["OCEAN_STATUS_TERMS"],
    }
    exec(
        compile(ast.Module(body=functions, type_ignores=[]), str(source), "exec"),
        namespace,
    )
    return namespace


helpers = load_helpers()


class DataRoutingTests(unittest.TestCase):
    def test_tuna_commissions_route_to_fisheries(self):
        route = helpers["_is_fisheries_prompt"]
        for source in ("IATTC", "ICCAT", "IOTC", "CCSBT", "EFFDIS_LL2000-2024"):
            self.assertTrue(route(f"查询 {source} 当前已入库数据"), source)

    def test_variable_only_v10_routes_to_ocean(self):
        route = helpers["_needs_ocean_mcp"]
        self.assertTrue(route("导出1998年1月4日135E-140E的v10数据"))
        self.assertFalse(route("provide a normal conversation response"))

    def test_variable_only_v10_gets_era5_and_netcdf_defaults(self):
        routed = helpers["_apply_ocean_export_defaults"](
            "导出1998年1月4日135E-140E的v10数据"
        )
        self.assertIn("source=era5", routed)
        self.assertIn("format=netcdf", routed)
        self.assertIn("mcp_marine_marine_export", routed)

    def test_export_without_real_tool_completion_is_rejected(self):
        guard = helpers["_ocean_export_execution_error"]
        error = guard(
            "导出1998年1月4日的v10数据",
            export_tool_completed=False,
            tool_result_text="",
            final_answer="正在向 Ocean 服务器提交，请稍后查询状态。",
        )
        self.assertIn("未实际执行", error)

    def test_export_with_real_download_url_is_accepted(self):
        guard = helpers["_ocean_export_execution_error"]
        error = guard(
            "导出1998年1月4日ERA5 v10数据",
            export_tool_completed=True,
            tool_result_text=(
                '{"status":"ok","download_url":'
                '"https://ocean.example/download/abc-123"}'
            ),
            final_answer="导出完成。",
        )
        self.assertEqual(error, "")


if __name__ == "__main__":
    unittest.main()