| import ast |
| import re |
| import unittest |
| from pathlib import Path |
|
|
|
|
| def load_helpers(): |
| source = Path(__file__).resolve().parents[1] / "ui_server.py" |
| tree = ast.parse(source.read_text(encoding="utf-8")) |
| names = { |
| "_is_confirmation_prompt", |
| "_is_ocean_export_request", |
| "_apply_ocean_export_defaults", |
| } |
| functions = [ |
| node for node in tree.body |
| if isinstance(node, ast.FunctionDef) and node.name in names |
| ] |
| namespace = { |
| "re": re, |
| "_needs_ocean_mcp": lambda prompt: any( |
| term in str(prompt).lower() |
| for term in ("era5", "ocean", "sst", "cmems", "oisst") |
| ), |
| } |
| exec( |
| compile(ast.Module(body=functions, type_ignores=[]), str(source), "exec"), |
| namespace, |
| ) |
| return namespace |
|
|
|
|
| helpers = load_helpers() |
|
|
|
|
| class FollowupRoutingTests(unittest.TestCase): |
| def test_short_confirmation_variants(self): |
| confirm = helpers["_is_confirmation_prompt"] |
| for value in ("确认", "好的", "继续", "OK", "yes"): |
| self.assertTrue(confirm(value)) |
| self.assertFalse(confirm("确认一下ERA5数据范围")) |
|
|
| def test_missing_format_defaults_to_netcdf(self): |
| apply_defaults = helpers["_apply_ocean_export_defaults"] |
| prompt = "导出1998年1月4日ERA5 v10数据" |
| routed = apply_defaults(prompt) |
| self.assertIn("format=netcdf", routed) |
| self.assertIn("mcp_marine_marine_export", routed) |
|
|
| def test_explicit_format_is_preserved(self): |
| apply_defaults = helpers["_apply_ocean_export_defaults"] |
| prompt = "导出1998年1月4日ERA5 v10数据，格式csv" |
| routed = apply_defaults(prompt) |
| self.assertIn(prompt, routed) |
| self.assertNotIn("format=netcdf", routed) |
| self.assertIn("mcp_marine_marine_export", routed) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
