1582 lines
59 KiB
Python
1582 lines
59 KiB
Python
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import socket
|
||
import sys
|
||
import tempfile
|
||
import time
|
||
import unittest
|
||
import uuid
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from copy import deepcopy
|
||
from pathlib import Path
|
||
from threading import Event, Lock, Thread
|
||
from types import SimpleNamespace
|
||
from unittest.mock import patch
|
||
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||
|
||
|
||
class ExternalCredentialCryptoTests(unittest.TestCase):
|
||
def test_requires_master_key_and_never_falls_back_to_plaintext(self):
|
||
from core.external_systems.crypto import encrypt_secret
|
||
|
||
with patch.dict(os.environ, {}, clear=True):
|
||
with self.assertRaisesRegex(RuntimeError, "ZCBOT_CREDENTIAL_MASTER_KEY"):
|
||
encrypt_secret("secret")
|
||
|
||
def test_roundtrip_uses_ciphertext(self):
|
||
from core.external_systems.crypto import decrypt_secret, encrypt_secret
|
||
|
||
with patch.dict(
|
||
os.environ,
|
||
{"ZCBOT_CREDENTIAL_MASTER_KEY": "unit-test-key-at-least-32-characters"},
|
||
clear=False,
|
||
):
|
||
stored = encrypt_secret("mes-password")
|
||
self.assertTrue(stored.startswith("v2:primary:"))
|
||
self.assertNotIn("mes-password", stored)
|
||
self.assertEqual(decrypt_secret(stored), "mes-password")
|
||
|
||
def test_ciphertext_is_bound_to_context_and_supports_key_rotation(self):
|
||
from core.external_systems.crypto import decrypt_secret, encrypt_secret
|
||
|
||
old_key = "old-unit-test-key-at-least-32-characters"
|
||
new_key = "new-unit-test-key-at-least-32-characters"
|
||
with patch.dict(
|
||
os.environ,
|
||
{
|
||
"ZCBOT_CREDENTIAL_MASTER_KEY": old_key,
|
||
"ZCBOT_CREDENTIAL_KEY_ID": "old",
|
||
},
|
||
clear=False,
|
||
):
|
||
stored = encrypt_secret("secret", aad="user:def:password")
|
||
with patch.dict(
|
||
os.environ,
|
||
{
|
||
"ZCBOT_CREDENTIAL_MASTER_KEY": new_key,
|
||
"ZCBOT_CREDENTIAL_KEY_ID": "new",
|
||
"ZCBOT_CREDENTIAL_PREVIOUS_KEYS": json.dumps({"old": old_key}),
|
||
},
|
||
clear=False,
|
||
):
|
||
self.assertEqual(
|
||
decrypt_secret(stored, aad="user:def:password"),
|
||
"secret",
|
||
)
|
||
with self.assertRaisesRegex(RuntimeError, "绑定上下文"):
|
||
decrypt_secret(stored, aad="other:def:password")
|
||
|
||
def test_rejects_short_master_key(self):
|
||
from core.external_systems.crypto import configured, encrypt_secret
|
||
|
||
with patch.dict(
|
||
os.environ, {"ZCBOT_CREDENTIAL_MASTER_KEY": "too-short"}, clear=False
|
||
):
|
||
self.assertFalse(configured())
|
||
with self.assertRaisesRegex(RuntimeError, "至少需要 32"):
|
||
encrypt_secret("mes-password")
|
||
|
||
|
||
class ExternalConnectionRevisionTests(unittest.TestCase):
|
||
def test_sensitive_definition_change_keeps_credentials_for_reverify(self):
|
||
from core.external_systems.service import _mark_connections_for_reverify
|
||
|
||
connection = SimpleNamespace(
|
||
credentials={"token": "ciphertext"},
|
||
credential_hint="***",
|
||
status="active",
|
||
last_verified_at="old",
|
||
last_error=None,
|
||
)
|
||
_mark_connections_for_reverify([connection])
|
||
self.assertEqual(connection.credentials, {"token": "ciphertext"})
|
||
self.assertEqual(connection.status, "needs_reverify")
|
||
self.assertIsNone(connection.last_verified_at)
|
||
|
||
def test_missing_credentials_remains_needs_credentials(self):
|
||
from core.external_systems.service import _mark_connections_for_reverify
|
||
|
||
connection = SimpleNamespace(
|
||
credentials={},
|
||
credential_hint="***",
|
||
status="active",
|
||
last_verified_at="old",
|
||
last_error=None,
|
||
)
|
||
_mark_connections_for_reverify([connection])
|
||
self.assertEqual(connection.credentials, {})
|
||
self.assertEqual(connection.status, "needs_credentials")
|
||
|
||
def test_legacy_missing_auth_defaults_are_semantically_unchanged(self):
|
||
from core.external_systems.service import _classify_definition_config_change
|
||
|
||
legacy = {
|
||
"base_url": "https://factory.invalid",
|
||
"openapi_url": "https://factory.invalid/swagger.json",
|
||
"login_path": "/api/auth/token/",
|
||
}
|
||
materialized = {
|
||
**legacy,
|
||
"auth_type": "password_jwt",
|
||
"username_field": "username",
|
||
"password_field": "password",
|
||
"token_field": "access",
|
||
"auth_header_name": "Authorization",
|
||
"auth_header_template": "Bearer {token}",
|
||
}
|
||
_, _, changed, impact = _classify_definition_config_change(
|
||
"factory_mes", legacy, materialized
|
||
)
|
||
self.assertEqual(changed, frozenset())
|
||
self.assertEqual(impact, "none")
|
||
|
||
def test_query_guidance_change_is_runtime_only(self):
|
||
from core.external_systems.service import _classify_definition_config_change
|
||
|
||
config = {
|
||
"base_url": "https://factory.invalid",
|
||
"openapi_url": "https://factory.invalid/swagger.json",
|
||
}
|
||
_, _, changed, impact = _classify_definition_config_change(
|
||
"factory_mes",
|
||
config,
|
||
{**config, "query_guidance": "先查数据集目录"},
|
||
)
|
||
self.assertEqual(changed, frozenset({"query_guidance"}))
|
||
self.assertEqual(impact, "runtime")
|
||
|
||
def test_target_change_requires_reverify_without_credential_reset(self):
|
||
from core.external_systems.service import _classify_definition_config_change
|
||
|
||
old = {
|
||
"base_url": "https://factory.invalid",
|
||
"openapi_url": "https://factory.invalid/swagger.json",
|
||
}
|
||
new = {
|
||
"base_url": "https://factory-new.invalid",
|
||
"openapi_url": "https://factory-new.invalid/swagger.json",
|
||
}
|
||
_, _, changed, impact = _classify_definition_config_change(
|
||
"factory_mes", old, new
|
||
)
|
||
self.assertEqual(
|
||
changed, frozenset({"base_url", "openapi_url"})
|
||
)
|
||
self.assertEqual(impact, "reverify")
|
||
|
||
def test_runtime_change_advances_only_active_connections(self):
|
||
from core.external_systems.service import _advance_active_connections
|
||
|
||
active = SimpleNamespace(
|
||
credentials={"token": "ciphertext"},
|
||
status="active",
|
||
verified_revision=1,
|
||
)
|
||
invalid = SimpleNamespace(
|
||
credentials={"token": "ciphertext"},
|
||
status="invalid",
|
||
verified_revision=1,
|
||
)
|
||
missing = SimpleNamespace(
|
||
credentials={},
|
||
status="active",
|
||
verified_revision=1,
|
||
last_verified_at="old",
|
||
last_error=None,
|
||
)
|
||
_advance_active_connections([active, invalid, missing], revision=2)
|
||
self.assertEqual(active.verified_revision, 2)
|
||
self.assertEqual(invalid.verified_revision, 1)
|
||
self.assertEqual(missing.status, "needs_credentials")
|
||
|
||
|
||
class GenericMcpConnectorTests(unittest.TestCase):
|
||
def _config(self, **overrides):
|
||
from core.external_systems.mcp import McpConfig
|
||
|
||
values = {
|
||
"base_url": "https://factory.invalid",
|
||
"mcp_url": "https://factory.invalid/mcp",
|
||
"auth_type": "api_key",
|
||
"auth_header_name": "Authorization",
|
||
"auth_header_template": "Bearer {token}",
|
||
"recommended_operation_ids": ["mcp/search_datasets"],
|
||
}
|
||
values.update(overrides)
|
||
return McpConfig.from_mapping(values)
|
||
|
||
def test_provider_uses_generic_mcp_without_tool_allowlist(self):
|
||
from core.external_systems.registry import get_provider
|
||
from core.external_systems.service import _normalized_config
|
||
|
||
provider = get_provider("generic_mcp")
|
||
normalized = _normalized_config(
|
||
"generic_mcp",
|
||
{
|
||
"mcp_url": "https://factory.invalid/mcp",
|
||
"auth_type": "api_key",
|
||
},
|
||
)
|
||
|
||
self.assertEqual(provider.connector, "mcp")
|
||
self.assertEqual(normalized["mcp_url"], "https://factory.invalid/mcp")
|
||
self.assertEqual(normalized["operation_policies"], {})
|
||
self.assertNotIn("openapi_url", normalized)
|
||
|
||
def test_mcp_url_and_login_origin_must_match(self):
|
||
from core.external_systems.mcp import McpConfig, McpConnectorError
|
||
|
||
with self.assertRaisesRegex(McpConnectorError, "同源"):
|
||
McpConfig.from_mapping(
|
||
{
|
||
"base_url": "https://login.invalid",
|
||
"mcp_url": "https://mcp.invalid/mcp",
|
||
}
|
||
)
|
||
|
||
def test_search_discovers_every_remote_tool_and_prioritizes_recommended(self):
|
||
from core.external_systems.mcp import McpClient
|
||
|
||
client = McpClient({"api_key": "secret"}, self._config())
|
||
catalog = {
|
||
"server": {"name": "factory", "version": "1"},
|
||
"tools": [
|
||
{
|
||
"name": "delete_future_tool",
|
||
"description": "服务器后来新增的工具",
|
||
"inputSchema": {"type": "object"},
|
||
},
|
||
{
|
||
"name": "search_datasets",
|
||
"description": "搜索数据集目录",
|
||
"inputSchema": {
|
||
"type": "object",
|
||
"properties": {"query": {"type": "string"}},
|
||
},
|
||
},
|
||
],
|
||
}
|
||
with patch.object(client, "_catalog", return_value=catalog):
|
||
results = client.search("服务器后来新增")
|
||
recommended = client.search("数据集")
|
||
|
||
self.assertIn(
|
||
"mcp/delete_future_tool",
|
||
{item["operation_id"] for item in results},
|
||
)
|
||
self.assertEqual(recommended[0]["operation_id"], "mcp/search_datasets")
|
||
self.assertEqual(
|
||
recommended[0]["input_schema"]["properties"]["query"]["type"],
|
||
"string",
|
||
)
|
||
|
||
def test_call_uses_remote_name_and_structured_content(self):
|
||
from core.external_systems.mcp import McpClient
|
||
|
||
client = McpClient({"api_key": "secret"}, self._config())
|
||
catalog = {
|
||
"server": {"name": "factory", "version": "1"},
|
||
"tools": [{"name": "get_wpr", "inputSchema": {"type": "object"}}],
|
||
}
|
||
with (
|
||
patch.object(client, "_catalog", return_value=catalog),
|
||
patch.object(
|
||
client,
|
||
"_run",
|
||
return_value={"structuredContent": {"number": "WPR-001"}},
|
||
) as called,
|
||
):
|
||
result = client.call("mcp/get_wpr", {"identifier": "WPR-001"})
|
||
|
||
called.assert_called_once_with(
|
||
"call", ("get_wpr", {"identifier": "WPR-001"})
|
||
)
|
||
self.assertEqual(result["data"], {"number": "WPR-001"})
|
||
self.assertEqual(result["operation_id"], "mcp/get_wpr")
|
||
|
||
def test_call_rejects_openapi_body_and_redacts_remote_secret(self):
|
||
from core.external_systems.mcp import McpClient, McpConnectorError
|
||
|
||
client = McpClient({"api_key": "secret"}, self._config())
|
||
catalog = {
|
||
"server": {"name": "factory", "version": "1"},
|
||
"tools": [{"name": "run", "inputSchema": {"type": "object"}}],
|
||
}
|
||
with patch.object(client, "_catalog", return_value=catalog):
|
||
with self.assertRaisesRegex(McpConnectorError, "arguments"):
|
||
client.call("mcp/run", body={"query": {}})
|
||
with (
|
||
patch.object(
|
||
client,
|
||
"_run",
|
||
return_value={
|
||
"isError": True,
|
||
"content": [{"type": "text", "token": "must-not-leak"}],
|
||
},
|
||
),
|
||
self.assertRaises(McpConnectorError) as raised,
|
||
):
|
||
client.call("mcp/run")
|
||
|
||
self.assertIn("[REDACTED]", str(raised.exception))
|
||
self.assertNotIn("must-not-leak", str(raised.exception))
|
||
|
||
def test_definition_target_change_requires_reverify(self):
|
||
from core.external_systems.service import _classify_definition_config_change
|
||
|
||
old = {"mcp_url": "https://factory.invalid/mcp"}
|
||
new = {"mcp_url": "https://factory.invalid/mcp-v2"}
|
||
_, _, changed, impact = _classify_definition_config_change(
|
||
"generic_mcp", old, new
|
||
)
|
||
|
||
self.assertEqual(changed, frozenset({"mcp_url"}))
|
||
self.assertEqual(impact, "reverify")
|
||
|
||
def test_streamable_http_server_is_discovered_and_called_end_to_end(self):
|
||
import uvicorn
|
||
from mcp.server import MCPServer
|
||
from mcp.server.transport_security import TransportSecuritySettings
|
||
|
||
from core.external_systems.mcp import McpClient, McpConnectorError
|
||
|
||
with socket.socket() as probe_socket:
|
||
probe_socket.bind(("127.0.0.1", 0))
|
||
port = probe_socket.getsockname()[1]
|
||
|
||
server_impl = MCPServer(
|
||
name="test-mcp",
|
||
title="Test MCP",
|
||
description="zcbot connector integration test",
|
||
version="1.0",
|
||
)
|
||
|
||
@server_impl.tool()
|
||
def echo_material(name: str) -> dict[str, str]:
|
||
"""返回材料名称。"""
|
||
return {"name": name}
|
||
|
||
@server_impl.tool()
|
||
def oversized_result() -> dict[str, str]:
|
||
"""返回超过客户端安全边界的测试内容。"""
|
||
return {"data": "x" * 70000}
|
||
|
||
app = server_impl.streamable_http_app(
|
||
streamable_http_path="/mcp",
|
||
json_response=True,
|
||
transport_security=TransportSecuritySettings(
|
||
enable_dns_rebinding_protection=True,
|
||
allowed_hosts=[f"127.0.0.1:{port}"],
|
||
allowed_origins=[],
|
||
),
|
||
host="127.0.0.1",
|
||
)
|
||
server = uvicorn.Server(
|
||
uvicorn.Config(
|
||
app,
|
||
host="127.0.0.1",
|
||
port=port,
|
||
log_level="warning",
|
||
)
|
||
)
|
||
thread = Thread(target=server.run, daemon=True)
|
||
thread.start()
|
||
deadline = time.time() + 5
|
||
while not server.started and time.time() < deadline:
|
||
time.sleep(0.01)
|
||
self.assertTrue(server.started)
|
||
|
||
client = McpClient(
|
||
{"api_key": "test-key"},
|
||
self._config(
|
||
base_url=f"http://127.0.0.1:{port}",
|
||
mcp_url=f"http://127.0.0.1:{port}/mcp",
|
||
expected_server_name="test-mcp",
|
||
recommended_operation_ids=[],
|
||
),
|
||
)
|
||
try:
|
||
connection = client.test_connection()
|
||
found = client.search("材料")
|
||
result = client.call(
|
||
"mcp/echo_material",
|
||
{"name": "低碳水泥"},
|
||
)
|
||
limited_client = McpClient(
|
||
{"api_key": "test-key"},
|
||
self._config(
|
||
base_url=f"http://127.0.0.1:{port}",
|
||
mcp_url=f"http://127.0.0.1:{port}/mcp",
|
||
expected_server_name="test-mcp",
|
||
recommended_operation_ids=[],
|
||
max_response_bytes=65536,
|
||
),
|
||
)
|
||
with self.assertRaisesRegex(McpConnectorError, "安全下载上限"):
|
||
limited_client.call("mcp/oversized_result")
|
||
finally:
|
||
server.should_exit = True
|
||
thread.join(timeout=5)
|
||
|
||
self.assertEqual(connection["server"]["name"], "test-mcp")
|
||
self.assertEqual(connection["tool_count"], 2)
|
||
self.assertEqual(found[0]["operation_id"], "mcp/echo_material")
|
||
self.assertEqual(result["data"], {"name": "低碳水泥"})
|
||
|
||
|
||
def _cfg(*, allowed=frozenset(), recommended=(), operation_mode="query"):
|
||
from core.external_systems.factory import FactoryMesConfig
|
||
|
||
return FactoryMesConfig(
|
||
base_url="https://factory.invalid",
|
||
openapi_url="https://factory.invalid/swagger.json",
|
||
login_path="/api/auth/token/",
|
||
operation_policies={operation_id: "read" for operation_id in allowed},
|
||
timeout_seconds=5,
|
||
max_result_bytes=65536,
|
||
max_total_result_bytes=262144,
|
||
max_page_size=200,
|
||
verify_tls=True,
|
||
query_guidance="先查数据集目录",
|
||
recommended_operation_ids=tuple(recommended),
|
||
operation_mode=operation_mode,
|
||
)
|
||
|
||
|
||
_SPEC = {
|
||
"swagger": "2.0",
|
||
"definitions": {
|
||
"DataExec": {
|
||
"type": "object",
|
||
"required": ["query"],
|
||
"properties": {
|
||
"query": {"title": "查询字典参数", "type": "object"},
|
||
"is_test": {"type": "boolean", "default": False},
|
||
},
|
||
}
|
||
},
|
||
"paths": {
|
||
"/api/bi/dataset/": {
|
||
"get": {
|
||
"operationId": "bi_dataset_list",
|
||
"summary": "复杂统计查询的数据集目录",
|
||
"tags": ["BI", "数据集", "报表"],
|
||
"parameters": [],
|
||
}
|
||
},
|
||
"/api/qm/ftestwork/{batch}/": {
|
||
"get": {
|
||
"operationId": "qm_ftestwork_read",
|
||
"summary": "查询成品检验批次",
|
||
"tags": ["quality"],
|
||
"parameters": [
|
||
{"name": "batch", "in": "path", "required": True, "type": "string"},
|
||
{
|
||
"name": "page_size",
|
||
"in": "query",
|
||
"required": False,
|
||
"type": "integer",
|
||
},
|
||
],
|
||
}
|
||
},
|
||
"/api/bi/dataset/{code}/exec/": {
|
||
"post": {
|
||
"operationId": "bi_dataset_exec",
|
||
"summary": "执行只读数据集",
|
||
"parameters": [
|
||
{"name": "code", "in": "path", "required": True, "type": "string"},
|
||
{
|
||
"name": "payload",
|
||
"in": "body",
|
||
"required": True,
|
||
"schema": {"$ref": "#/definitions/DataExec"},
|
||
},
|
||
],
|
||
}
|
||
},
|
||
},
|
||
}
|
||
|
||
|
||
class _Response:
|
||
def __init__(self, status_code=200, payload=None):
|
||
self.status_code = status_code
|
||
self._payload = payload
|
||
self.headers = {"content-type": "application/json"}
|
||
self.text = json.dumps(payload, ensure_ascii=False)
|
||
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, *args):
|
||
return False
|
||
|
||
def iter_bytes(self):
|
||
yield self.text.encode("utf-8")
|
||
|
||
def json(self):
|
||
return self._payload
|
||
|
||
|
||
class _Http:
|
||
def __init__(self):
|
||
self.calls = []
|
||
|
||
def __enter__(self):
|
||
return self
|
||
|
||
def __exit__(self, *args):
|
||
return False
|
||
|
||
def post(self, url, **kwargs):
|
||
self.calls.append(("POST", url, kwargs))
|
||
return _Response(payload={"access": "remote-jwt"})
|
||
|
||
def get(self, url, **kwargs):
|
||
self.calls.append(("GET", url, kwargs))
|
||
return _Response(payload=_SPEC)
|
||
|
||
def request(self, method, url, **kwargs):
|
||
self.calls.append((method, url, kwargs))
|
||
return _Response(payload={"count": 1, "results": [{"batch": "B/1"}]})
|
||
|
||
def stream(self, method, url, **kwargs):
|
||
if method.upper() == "POST" and "/auth/" in url:
|
||
return self.post(url, **kwargs)
|
||
if method.upper() == "GET" and ("swagger" in url or "openapi" in url):
|
||
return self.get(url, **kwargs)
|
||
return self.request(method.upper(), url, **kwargs)
|
||
|
||
|
||
class ExternalRuntimeCacheTests(unittest.TestCase):
|
||
def test_lru_defers_client_close_until_active_lease_finishes(self):
|
||
from core.external_systems.runtime_cache import ExternalRuntimeCache
|
||
|
||
class Client:
|
||
def __init__(self):
|
||
self.closed = False
|
||
|
||
def close(self):
|
||
self.closed = True
|
||
|
||
cache = ExternalRuntimeCache(max_entries=2)
|
||
first = Client()
|
||
second = Client()
|
||
third = Client()
|
||
with cache.client("first", lambda: first):
|
||
with cache.client("second", lambda: second):
|
||
pass
|
||
with cache.client("third", lambda: third):
|
||
pass
|
||
self.assertFalse(first.closed)
|
||
self.assertTrue(first.closed)
|
||
cache.clear()
|
||
self.assertTrue(second.closed)
|
||
self.assertTrue(third.closed)
|
||
|
||
|
||
class FactoryOpenApiConnectorTests(unittest.TestCase):
|
||
def setUp(self):
|
||
from core.external_systems.openapi import _SPEC_CACHE
|
||
|
||
_SPEC_CACHE.clear()
|
||
|
||
def test_admin_mapping_builds_bounded_runtime_config(self):
|
||
from core.external_systems.factory import FactoryMesConfig
|
||
|
||
cfg = FactoryMesConfig.from_mapping(
|
||
{
|
||
"base_url": "https://factory.invalid/",
|
||
"openapi_url": "https://factory.invalid/swagger.json",
|
||
"operation_policies": {
|
||
"bi_dataset_exec": "read",
|
||
"report_preview": "export",
|
||
},
|
||
"timeout_seconds": 999,
|
||
"max_result_bytes": 1,
|
||
"verify_tls": True,
|
||
}
|
||
)
|
||
self.assertEqual(cfg.base_url, "https://factory.invalid")
|
||
self.assertEqual(cfg.timeout_seconds, 60)
|
||
self.assertEqual(cfg.max_result_bytes, 4096)
|
||
self.assertEqual(cfg.max_total_result_bytes, 262144)
|
||
self.assertEqual(cfg.max_page_size, 200)
|
||
self.assertEqual(cfg.operation_mode, "upstream_managed")
|
||
self.assertEqual(
|
||
cfg.operation_policies,
|
||
{
|
||
"bi_dataset_exec": "read",
|
||
"report_preview": "export",
|
||
},
|
||
)
|
||
self.assertIn("dataset list", cfg.query_guidance)
|
||
self.assertEqual(
|
||
cfg.recommended_operation_ids,
|
||
("bi_dataset_list", "bi_dataset_exec"),
|
||
)
|
||
|
||
def test_admin_mapping_rejects_embedded_url_credentials(self):
|
||
from core.external_systems.factory import FactoryMesConfig, FactoryMesError
|
||
|
||
with self.assertRaisesRegex(FactoryMesError, "不能内嵌凭据"):
|
||
FactoryMesConfig.from_mapping(
|
||
{
|
||
"base_url": "https://user:secret@factory.invalid",
|
||
"openapi_url": "https://factory.invalid/swagger.json",
|
||
}
|
||
)
|
||
|
||
def test_admin_mapping_requires_same_origin_openapi_document(self):
|
||
from core.external_systems.factory import FactoryMesConfig, FactoryMesError
|
||
|
||
with self.assertRaisesRegex(FactoryMesError, "必须与 base_url 同源"):
|
||
FactoryMesConfig.from_mapping(
|
||
{
|
||
"base_url": "https://factory.invalid",
|
||
"openapi_url": "https://spec.attacker.invalid/swagger.json",
|
||
}
|
||
)
|
||
|
||
def test_generic_api_key_auth_uses_declared_header_without_login(self):
|
||
from core.external_systems.openapi import (
|
||
_SPEC_CACHE,
|
||
OpenApiClient,
|
||
OpenApiConfig,
|
||
)
|
||
from core.external_systems.registry import credential_fields, merged_config
|
||
|
||
_SPEC_CACHE.clear()
|
||
config = merged_config(
|
||
"generic_openapi",
|
||
{
|
||
"base_url": "https://erp.invalid",
|
||
"openapi_url": "https://erp.invalid/openapi.json",
|
||
"auth_type": "api_key",
|
||
"auth_header_name": "X-ERP-Key",
|
||
"auth_header_template": "Key {token}",
|
||
},
|
||
)
|
||
client = OpenApiClient(
|
||
{"api_key": "private-key"}, OpenApiConfig.from_mapping(config)
|
||
)
|
||
http = _Http()
|
||
with patch.object(client, "_client", return_value=http):
|
||
result = client.test_connection()
|
||
self.assertGreater(result["operation_count"], 0)
|
||
self.assertFalse(any(call[0] == "POST" for call in http.calls))
|
||
get_call = next(call for call in http.calls if call[0] == "GET")
|
||
self.assertEqual(get_call[2]["headers"], {"X-ERP-Key": "Key private-key"})
|
||
self.assertEqual(
|
||
credential_fields("generic_openapi", config)[0]["name"], "api_key"
|
||
)
|
||
|
||
def test_openapi_spec_cache_is_isolated_by_connection_namespace(self):
|
||
from core.external_systems.openapi import (
|
||
_SPEC_CACHE,
|
||
OpenApiClient,
|
||
OpenApiConfig,
|
||
)
|
||
from core.external_systems.registry import merged_config
|
||
|
||
_SPEC_CACHE.clear()
|
||
config = OpenApiConfig.from_mapping(
|
||
merged_config(
|
||
"generic_openapi",
|
||
{
|
||
"base_url": "https://erp.invalid",
|
||
"openapi_url": "https://erp.invalid/openapi.json",
|
||
"auth_type": "bearer_token",
|
||
},
|
||
)
|
||
)
|
||
for namespace in ("definition:user-a", "definition:user-b"):
|
||
client = OpenApiClient(
|
||
{"token": namespace}, config, cache_namespace=namespace
|
||
)
|
||
http = _Http()
|
||
with patch.object(client, "_client", return_value=http):
|
||
client.test_connection()
|
||
self.assertTrue(any(call[0] == "GET" for call in http.calls))
|
||
self.assertEqual(len(_SPEC_CACHE), 2)
|
||
|
||
def test_search_discovers_operation_without_exposing_credentials(self):
|
||
from core.external_systems.factory import FactoryMesClient
|
||
|
||
http = _Http()
|
||
client = FactoryMesClient("mes-user", "mes-password", _cfg())
|
||
with patch.object(client, "_client", return_value=http):
|
||
result = client.search("成品检验")
|
||
self.assertEqual(result[0]["operation_id"], "qm_ftestwork_read")
|
||
rendered = json.dumps(result, ensure_ascii=False)
|
||
self.assertNotIn("mes-password", rendered)
|
||
self.assertNotIn("remote-jwt", rendered)
|
||
|
||
def test_runtime_reuses_http_auth_spec_and_compiled_catalog(self):
|
||
from core.external_systems.factory import FactoryMesClient
|
||
from core.external_systems.openapi import compile_operation_catalog
|
||
|
||
http = _Http()
|
||
first = FactoryMesClient("mes-user", "mes-password", _cfg())
|
||
second = FactoryMesClient("mes-user", "mes-password", _cfg())
|
||
with (
|
||
patch.object(first, "_client", return_value=http) as first_factory,
|
||
patch.object(second, "_client", return_value=_Http()) as second_factory,
|
||
patch(
|
||
"core.external_systems.openapi.compile_operation_catalog",
|
||
wraps=compile_operation_catalog,
|
||
) as compile_catalog,
|
||
):
|
||
first.search("成品检验")
|
||
second.search("成品检验")
|
||
first.call("qm_ftestwork_read", arguments={"batch": "B1"})
|
||
second.call("qm_ftestwork_read", arguments={"batch": "B1"})
|
||
|
||
self.assertEqual(first_factory.call_count, 1)
|
||
self.assertEqual(second_factory.call_count, 0)
|
||
self.assertEqual(compile_catalog.call_count, 1)
|
||
self.assertEqual(
|
||
sum(
|
||
1
|
||
for method, url, _ in http.calls
|
||
if method == "POST" and "/auth/" in url
|
||
),
|
||
1,
|
||
)
|
||
self.assertEqual(
|
||
sum(
|
||
1
|
||
for method, url, _ in http.calls
|
||
if method == "GET" and "swagger" in url
|
||
),
|
||
1,
|
||
)
|
||
self.assertEqual(
|
||
sum(
|
||
1
|
||
for method, url, _ in http.calls
|
||
if method == "GET" and "/ftestwork/" in url
|
||
),
|
||
2,
|
||
)
|
||
|
||
def test_concurrent_identical_query_is_singleflight_only(self):
|
||
from core.external_systems.factory import FactoryMesClient
|
||
|
||
class SlowHttp(_Http):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.query_started = Event()
|
||
self.release_query = Event()
|
||
self.query_count = 0
|
||
self.query_lock = Lock()
|
||
|
||
def request(self, method, url, **kwargs):
|
||
if method == "GET" and "/ftestwork/" in url:
|
||
with self.query_lock:
|
||
self.query_count += 1
|
||
self.query_started.set()
|
||
self.release_query.wait(timeout=2)
|
||
return super().request(method, url, **kwargs)
|
||
|
||
http = SlowHttp()
|
||
client = FactoryMesClient("mes-user", "mes-password", _cfg())
|
||
with patch.object(client, "_client", return_value=http):
|
||
client.search("成品检验") # 预热认证、规格和 catalog，只测业务请求单飞。
|
||
with ThreadPoolExecutor(max_workers=2) as executor:
|
||
first = executor.submit(
|
||
client.call,
|
||
"qm_ftestwork_read",
|
||
{"batch": "B1"},
|
||
)
|
||
self.assertTrue(http.query_started.wait(timeout=1))
|
||
second = executor.submit(
|
||
client.call,
|
||
"qm_ftestwork_read",
|
||
{"batch": "B1"},
|
||
)
|
||
time.sleep(0.05)
|
||
http.release_query.set()
|
||
first_result = first.result(timeout=2)
|
||
second_result = second.result(timeout=2)
|
||
|
||
self.assertEqual(http.query_count, 1)
|
||
self.assertEqual(first_result, second_result)
|
||
self.assertIsNot(first_result, second_result)
|
||
|
||
def test_concurrent_cold_search_coalesces_login_and_spec_fetch(self):
|
||
from core.external_systems.factory import FactoryMesClient
|
||
|
||
class SlowDiscoveryHttp(_Http):
|
||
def post(self, url, **kwargs):
|
||
response = super().post(url, **kwargs)
|
||
time.sleep(0.05)
|
||
return response
|
||
|
||
def get(self, url, **kwargs):
|
||
response = super().get(url, **kwargs)
|
||
time.sleep(0.05)
|
||
return response
|
||
|
||
http = SlowDiscoveryHttp()
|
||
first = FactoryMesClient("mes-user", "mes-password", _cfg())
|
||
second = FactoryMesClient("mes-user", "mes-password", _cfg())
|
||
with (
|
||
patch.object(first, "_client", return_value=http),
|
||
patch.object(second, "_client", return_value=http),
|
||
ThreadPoolExecutor(max_workers=2) as executor,
|
||
):
|
||
results = list(
|
||
executor.map(lambda client: client.search("成品检验"), (first, second))
|
||
)
|
||
|
||
self.assertTrue(all(result for result in results))
|
||
self.assertEqual(
|
||
sum(
|
||
1
|
||
for method, url, _ in http.calls
|
||
if method == "POST" and "/auth/" in url
|
||
),
|
||
1,
|
||
)
|
||
self.assertEqual(
|
||
sum(
|
||
1
|
||
for method, url, _ in http.calls
|
||
if method == "GET" and "swagger" in url
|
||
),
|
||
1,
|
||
)
|
||
|
||
def test_cached_password_token_refreshes_once_after_401(self):
|
||
from core.external_systems.factory import FactoryMesClient
|
||
|
||
class RefreshHttp(_Http):
|
||
def __init__(self):
|
||
super().__init__()
|
||
self.business_attempts = 0
|
||
|
||
def request(self, method, url, **kwargs):
|
||
if method == "GET" and "/ftestwork/" in url:
|
||
self.calls.append((method, url, kwargs))
|
||
self.business_attempts += 1
|
||
if self.business_attempts == 1:
|
||
return _Response(payload={"detail": "expired"}, status_code=401)
|
||
return super().request(method, url, **kwargs)
|
||
|
||
http = RefreshHttp()
|
||
client = FactoryMesClient("mes-user", "mes-password", _cfg())
|
||
with patch.object(client, "_client", return_value=http):
|
||
result = client.call("qm_ftestwork_read", arguments={"batch": "B1"})
|
||
|
||
self.assertEqual(result["status_code"], 200)
|
||
self.assertEqual(http.business_attempts, 2)
|
||
self.assertEqual(
|
||
sum(
|
||
1
|
||
for method, url, _ in http.calls
|
||
if method == "POST" and "/auth/" in url
|
||
),
|
||
2,
|
||
)
|
||
|
||
def test_search_pins_callable_admin_recommendations_without_keyword_match(self):
|
||
from core.external_systems.factory import FactoryMesClient
|
||
|
||
http = _Http()
|
||
client = FactoryMesClient(
|
||
"mes-user",
|
||
"mes-password",
|
||
_cfg(
|
||
allowed={"bi_dataset_exec"},
|
||
recommended=("bi_dataset_list", "bi_dataset_exec"),
|
||
),
|
||
)
|
||
with patch.object(client, "_client", return_value=http):
|
||
result = client.search("某工段上月产量")
|
||
self.assertEqual(
|
||
[item["operation_id"] for item in result[:2]],
|
||
["bi_dataset_list", "bi_dataset_exec"],
|
||
)
|
||
self.assertTrue(all(item["recommended"] for item in result[:2]))
|
||
|
||
def test_search_includes_resolved_swagger_body_schema(self):
|
||
from core.external_systems.factory import FactoryMesClient
|
||
|
||
http = _Http()
|
||
client = FactoryMesClient(
|
||
"mes-user", "mes-password", _cfg(allowed={"bi_dataset_exec"})
|
||
)
|
||
with patch.object(client, "_client", return_value=http):
|
||
result = client.search("执行只读数据集")
|
||
operation = next(
|
||
item for item in result if item["operation_id"] == "bi_dataset_exec"
|
||
)
|
||
self.assertEqual(operation["body"]["parameter_name"], "payload")
|
||
self.assertTrue(operation["body"]["required"])
|
||
self.assertEqual(
|
||
operation["body"]["schema"]["properties"]["query"]["type"],
|
||
"object",
|
||
)
|
||
self.assertEqual(operation["body"]["schema"]["required"], ["query"])
|
||
|
||
def test_search_includes_openapi_3_request_body_schema(self):
|
||
from core.external_systems.openapi import OpenApiClient
|
||
|
||
spec = {
|
||
"openapi": "3.0.0",
|
||
"components": {
|
||
"schemas": {
|
||
"ReportQuery": {
|
||
"type": "object",
|
||
"properties": {"month": {"type": "string"}},
|
||
}
|
||
}
|
||
},
|
||
"paths": {
|
||
"/reports/preview/": {
|
||
"post": {
|
||
"operationId": "report_preview",
|
||
"summary": "预览报表",
|
||
"requestBody": {
|
||
"required": True,
|
||
"content": {
|
||
"application/json": {
|
||
"schema": {
|
||
"$ref": "#/components/schemas/ReportQuery"
|
||
}
|
||
}
|
||
},
|
||
},
|
||
}
|
||
}
|
||
},
|
||
}
|
||
client = OpenApiClient({"token": "private"}, _cfg(allowed={"report_preview"}))
|
||
with (
|
||
patch.object(client, "authenticate", return_value={}),
|
||
patch.object(client, "_fetch_spec", return_value=spec),
|
||
):
|
||
result = client.search("预览报表")
|
||
self.assertEqual(result[0]["body"]["content_type"], "application/json")
|
||
self.assertEqual(
|
||
result[0]["body"]["schema"]["properties"]["month"]["type"],
|
||
"string",
|
||
)
|
||
|
||
def test_catalog_resolves_referenced_header_parameter(self):
|
||
from core.external_systems.factory import FactoryMesClient
|
||
|
||
spec = {
|
||
"openapi": "3.0.0",
|
||
"components": {
|
||
"parameters": {
|
||
"Trace": {
|
||
"name": "X-Trace-Id",
|
||
"in": "header",
|
||
"required": True,
|
||
"schema": {"type": "string"},
|
||
}
|
||
}
|
||
},
|
||
"paths": {
|
||
"/quality/": {
|
||
"get": {
|
||
"operationId": "quality_read",
|
||
"parameters": [{"$ref": "#/components/parameters/Trace"}],
|
||
}
|
||
}
|
||
},
|
||
}
|
||
cfg = _cfg()
|
||
http = _Http()
|
||
client = FactoryMesClient("u", "p", cfg)
|
||
with (
|
||
patch.object(client, "_client", return_value=http),
|
||
patch.object(client, "_fetch_spec", return_value=spec),
|
||
):
|
||
client.call("quality_read", arguments={"X-Trace-Id": "trace-1"})
|
||
request = next(call for call in http.calls if call[0] == "GET")
|
||
self.assertEqual(request[2]["headers"]["X-Trace-Id"], "trace-1")
|
||
|
||
def test_swagger_array_query_uses_declared_collection_format(self):
|
||
from core.external_systems.factory import FactoryMesClient
|
||
|
||
spec = {
|
||
"swagger": "2.0",
|
||
"paths": {
|
||
"/quality/": {
|
||
"get": {
|
||
"operationId": "quality_filter",
|
||
"parameters": [
|
||
{
|
||
"name": "batches",
|
||
"in": "query",
|
||
"type": "array",
|
||
"items": {"type": "string"},
|
||
"collectionFormat": "csv",
|
||
}
|
||
],
|
||
}
|
||
}
|
||
},
|
||
}
|
||
cfg = _cfg()
|
||
http = _Http()
|
||
client = FactoryMesClient("u", "p", cfg)
|
||
with (
|
||
patch.object(client, "_client", return_value=http),
|
||
patch.object(client, "_fetch_spec", return_value=spec),
|
||
):
|
||
client.call("quality_filter", arguments={"batches": ["B1", "B2"]})
|
||
request = next(call for call in http.calls if call[0] == "GET")
|
||
self.assertEqual(request[2]["params"]["batches"], "B1,B2")
|
||
|
||
def test_get_call_resolves_encoded_path_and_query(self):
|
||
from core.external_systems.factory import FactoryMesClient
|
||
|
||
http = _Http()
|
||
client = FactoryMesClient("mes-user", "mes-password", _cfg())
|
||
with patch.object(client, "_client", return_value=http):
|
||
result = client.call(
|
||
"qm_ftestwork_read",
|
||
arguments={"batch": "B/1", "page_size": 50},
|
||
)
|
||
method, url, kwargs = [
|
||
call for call in http.calls if call[0] == "GET" and "/api/" in call[1]
|
||
][0]
|
||
self.assertEqual(method, "GET")
|
||
self.assertIn("B%2F1", url)
|
||
self.assertEqual(kwargs["params"], {"page_size": 50})
|
||
self.assertEqual(result["data"]["count"], 1)
|
||
|
||
def test_get_call_bounds_pagination_for_agent_queries(self):
|
||
from core.external_systems.factory import FactoryMesClient, FactoryMesError
|
||
|
||
spec = deepcopy(_SPEC)
|
||
spec["paths"]["/api/qm/ftestwork/{batch}/"]["get"]["parameters"].extend(
|
||
[
|
||
{"name": "page", "in": "query", "required": False, "type": "integer"},
|
||
{
|
||
"name": "pageoff",
|
||
"in": "query",
|
||
"required": False,
|
||
"type": "boolean",
|
||
},
|
||
]
|
||
)
|
||
http = _Http()
|
||
client = FactoryMesClient("u", "p", _cfg())
|
||
with (
|
||
patch.object(client, "_client", return_value=http),
|
||
patch.object(client, "_fetch_spec", return_value=spec),
|
||
):
|
||
client.call(
|
||
"qm_ftestwork_read",
|
||
arguments={"batch": "B1", "page": 1, "page_size": 99999},
|
||
)
|
||
request = next(call for call in http.calls if call[0] == "GET")
|
||
self.assertEqual(request[2]["params"]["page_size"], 200)
|
||
|
||
with (
|
||
patch.object(client, "authenticate", return_value="jwt"),
|
||
patch.object(client, "_fetch_spec", return_value=spec),
|
||
self.assertRaisesRegex(FactoryMesError, "不允许 page=0"),
|
||
):
|
||
client.call(
|
||
"qm_ftestwork_read",
|
||
arguments={"batch": "B1", "page": 0},
|
||
)
|
||
|
||
def test_swagger_base_path_is_added_to_operation_url(self):
|
||
from core.external_systems.factory import FactoryMesClient
|
||
|
||
spec = deepcopy(_SPEC)
|
||
spec["basePath"] = "/api"
|
||
spec["paths"] = {
|
||
path.removeprefix("/api"): value for path, value in spec["paths"].items()
|
||
}
|
||
http = _Http()
|
||
client = FactoryMesClient("u", "p", _cfg())
|
||
with (
|
||
patch.object(client, "_client", return_value=http),
|
||
patch.object(client, "_fetch_spec", return_value=spec),
|
||
):
|
||
client.call("qm_ftestwork_read", arguments={"batch": "B1"})
|
||
request = next(call for call in http.calls if call[0] == "GET")
|
||
self.assertEqual(request[1], "https://factory.invalid/api/qm/ftestwork/B1/")
|
||
|
||
def test_api_base_path_does_not_change_login_url(self):
|
||
from core.external_systems.factory import FactoryMesClient
|
||
|
||
http = _Http()
|
||
client = FactoryMesClient("u", "p", _cfg())
|
||
with patch.object(client, "_client", return_value=http):
|
||
client.authenticate()
|
||
request = next(call for call in http.calls if call[0] == "POST")
|
||
self.assertEqual(request[1], "https://factory.invalid/api/auth/token/")
|
||
|
||
def test_base_path_is_not_duplicated_when_operation_already_contains_it(self):
|
||
from core.external_systems.factory import FactoryMesClient, FactoryMesConfig
|
||
|
||
spec = {**deepcopy(_SPEC), "basePath": "/api"}
|
||
http = _Http()
|
||
client = FactoryMesClient("u", "p", _cfg())
|
||
with (
|
||
patch.object(client, "_client", return_value=http),
|
||
patch.object(client, "_fetch_spec", return_value=spec),
|
||
):
|
||
client.call("qm_ftestwork_read", arguments={"batch": "B1"})
|
||
request = next(call for call in http.calls if call[0] == "GET")
|
||
self.assertNotIn("/api/api/", request[1])
|
||
|
||
configured_prefix = FactoryMesClient(
|
||
"u",
|
||
"p",
|
||
FactoryMesConfig(
|
||
**{**_cfg().__dict__, "base_url": "https://factory.invalid/api"}
|
||
),
|
||
)
|
||
self.assertEqual(
|
||
configured_prefix._operation_url(spec, "/api/qm/ftestwork/B1/"),
|
||
"https://factory.invalid/api/qm/ftestwork/B1/",
|
||
)
|
||
|
||
def test_openapi_server_path_is_used_but_cross_origin_server_is_rejected(self):
|
||
from core.external_systems.factory import FactoryMesClient, FactoryMesError
|
||
|
||
client = FactoryMesClient("u", "p", _cfg())
|
||
same_origin = {"openapi": "3.0.0", "servers": [{"url": "/v1"}], "paths": {}}
|
||
self.assertEqual(
|
||
client._operation_url(same_origin, "/quality/results/"),
|
||
"https://factory.invalid/v1/quality/results/",
|
||
)
|
||
cross_origin = {
|
||
"openapi": "3.0.0",
|
||
"servers": [{"url": "https://attacker.invalid/v1"}],
|
||
"paths": {},
|
||
}
|
||
with self.assertRaisesRegex(FactoryMesError, "越出 Factory MES 主机"):
|
||
client._operation_url(cross_origin, "/quality/results/")
|
||
|
||
def test_no_declared_base_path_keeps_existing_url_behavior(self):
|
||
from core.external_systems.factory import FactoryMesClient
|
||
|
||
client = FactoryMesClient("u", "p", _cfg())
|
||
self.assertEqual(
|
||
client._operation_url(_SPEC, "/api/qm/ftestwork/B1/"),
|
||
"https://factory.invalid/api/qm/ftestwork/B1/",
|
||
)
|
||
|
||
def test_post_is_denied_unless_admin_allowlists_operation(self):
|
||
from core.external_systems.factory import FactoryMesClient, FactoryMesError
|
||
|
||
denied = FactoryMesClient("u", "p", _cfg())
|
||
with (
|
||
patch.object(denied, "authenticate", return_value="jwt"),
|
||
patch.object(denied, "_fetch_spec", return_value=_SPEC),
|
||
):
|
||
with self.assertRaisesRegex(FactoryMesError, "只读调用范围"):
|
||
denied.call("bi_dataset_exec", arguments={"code": "x", "payload": {}})
|
||
|
||
http = _Http()
|
||
allowed = FactoryMesClient("u", "p", _cfg(allowed={"bi_dataset_exec"}))
|
||
with patch.object(allowed, "_client", return_value=http):
|
||
result = allowed.call(
|
||
"bi_dataset_exec",
|
||
arguments={"code": "yield", "payload": {"query": {"month": "2026-08"}}},
|
||
)
|
||
request = [
|
||
call for call in http.calls if call[0] == "POST" and "/dataset/" in call[1]
|
||
][0]
|
||
self.assertEqual(request[2]["json"], {"query": {"month": "2026-08"}})
|
||
self.assertFalse(result["truncated"])
|
||
|
||
def test_allowlisted_post_accepts_separate_body_field(self):
|
||
from core.external_systems.factory import FactoryMesClient
|
||
|
||
http = _Http()
|
||
allowed = FactoryMesClient("u", "p", _cfg(allowed={"bi_dataset_exec"}))
|
||
with patch.object(allowed, "_client", return_value=http):
|
||
result = allowed.call(
|
||
"bi_dataset_exec",
|
||
arguments={"code": "quality"},
|
||
body={"query": {"batch": "B-1"}},
|
||
)
|
||
request = [
|
||
call for call in http.calls if call[0] == "POST" and "/dataset/" in call[1]
|
||
][0]
|
||
self.assertEqual(request[2]["json"], {"query": {"batch": "B-1"}})
|
||
self.assertFalse(result["truncated"])
|
||
|
||
def test_upstream_managed_mode_allows_declared_write_method(self):
|
||
from core.external_systems.factory import FactoryMesClient
|
||
|
||
spec = deepcopy(_SPEC)
|
||
spec["paths"]["/api/items/{item_id}/"] = {
|
||
"put": {
|
||
"operationId": "item_update",
|
||
"parameters": [
|
||
{
|
||
"name": "item_id",
|
||
"in": "path",
|
||
"required": True,
|
||
"type": "string",
|
||
},
|
||
{
|
||
"name": "payload",
|
||
"in": "body",
|
||
"required": True,
|
||
"schema": {
|
||
"type": "object",
|
||
"required": ["name"],
|
||
"properties": {"name": {"type": "string"}},
|
||
},
|
||
},
|
||
],
|
||
}
|
||
}
|
||
http = _Http()
|
||
client = FactoryMesClient("u", "p", _cfg(operation_mode="upstream_managed"))
|
||
with (
|
||
patch.object(client, "_client", return_value=http),
|
||
patch.object(client, "authenticate", return_value={}),
|
||
patch.object(client, "_fetch_spec", return_value=spec),
|
||
):
|
||
result = client.call(
|
||
"item_update",
|
||
arguments={"item_id": "A/B"},
|
||
body={"name": "updated"},
|
||
)
|
||
request = next(call for call in http.calls if call[0] == "PUT")
|
||
self.assertIn("/api/items/A%2FB/", request[1])
|
||
self.assertEqual(request[2]["json"], {"name": "updated"})
|
||
self.assertEqual(result["status_code"], 200)
|
||
|
||
def test_query_mode_still_rejects_declared_write_method(self):
|
||
from core.external_systems.factory import FactoryMesClient, FactoryMesError
|
||
|
||
spec = {
|
||
"swagger": "2.0",
|
||
"paths": {
|
||
"/api/items/{item_id}/": {
|
||
"delete": {
|
||
"operationId": "item_delete",
|
||
"parameters": [
|
||
{
|
||
"name": "item_id",
|
||
"in": "path",
|
||
"required": True,
|
||
"type": "string",
|
||
}
|
||
],
|
||
}
|
||
}
|
||
},
|
||
}
|
||
client = FactoryMesClient("u", "p", _cfg())
|
||
with (
|
||
patch.object(client, "authenticate", return_value={}),
|
||
patch.object(client, "_fetch_spec", return_value=spec),
|
||
):
|
||
with self.assertRaisesRegex(FactoryMesError, "只读调用范围"):
|
||
client.call("item_delete", arguments={"item_id": "A-1"})
|
||
|
||
def test_call_preserves_payload_larger_than_inline_limit(self):
|
||
from core.external_systems.factory import FactoryMesClient
|
||
|
||
class LargeHttp(_Http):
|
||
def request(self, method, url, **kwargs):
|
||
self.calls.append((method, url, kwargs))
|
||
return _Response(payload={"rows": "x" * 70000})
|
||
|
||
client = FactoryMesClient("u", "p", _cfg(allowed={"bi_dataset_exec"}))
|
||
with patch.object(client, "_client", return_value=LargeHttp()):
|
||
result = client.call(
|
||
"bi_dataset_exec",
|
||
arguments={"code": "quality"},
|
||
body={"query": {}},
|
||
)
|
||
self.assertEqual(len(result["data"]["rows"]), 70000)
|
||
self.assertGreater(result["response_bytes"], client.cfg.max_result_bytes)
|
||
self.assertFalse(result["truncated"])
|
||
|
||
def test_call_stops_stream_when_download_limit_is_exceeded(self):
|
||
from core.external_systems.factory import FactoryMesClient, FactoryMesError
|
||
from core.external_systems.results import MAX_STORED_RESULT_BYTES
|
||
|
||
class OversizedResponse(_Response):
|
||
def iter_bytes(self):
|
||
yield b"x" * MAX_STORED_RESULT_BYTES
|
||
yield b"x"
|
||
|
||
class OversizedHttp(_Http):
|
||
def request(self, method, url, **kwargs):
|
||
self.calls.append((method, url, kwargs))
|
||
return OversizedResponse(payload={})
|
||
|
||
client = FactoryMesClient("u", "p", _cfg())
|
||
with patch.object(client, "_client", return_value=OversizedHttp()):
|
||
with self.assertRaisesRegex(FactoryMesError, "安全下载上限"):
|
||
client.call("qm_ftestwork_read", arguments={"batch": "B1"})
|
||
|
||
def test_call_surfaces_sanitized_upstream_error_detail(self):
|
||
from core.external_systems.factory import FactoryMesClient, FactoryMesError
|
||
|
||
class ErrorHttp(_Http):
|
||
def request(self, method, url, **kwargs):
|
||
self.calls.append((method, url, kwargs))
|
||
return _Response(
|
||
status_code=400,
|
||
payload={
|
||
"query": ["This field is required."],
|
||
"access_token": "must-not-leak",
|
||
},
|
||
)
|
||
|
||
client = FactoryMesClient("u", "p", _cfg(allowed={"bi_dataset_exec"}))
|
||
with patch.object(client, "_client", return_value=ErrorHttp()):
|
||
with self.assertRaises(FactoryMesError) as raised:
|
||
client.call(
|
||
"bi_dataset_exec",
|
||
arguments={"code": "yield"},
|
||
body={"query": {}},
|
||
)
|
||
message = str(raised.exception)
|
||
self.assertIn("This field is required.", message)
|
||
self.assertIn("[REDACTED]", message)
|
||
self.assertNotIn("must-not-leak", message)
|
||
|
||
def test_rejects_unknown_arguments(self):
|
||
from core.external_systems.factory import FactoryMesClient, FactoryMesError
|
||
|
||
client = FactoryMesClient("u", "p", _cfg())
|
||
with (
|
||
patch.object(client, "authenticate", return_value="jwt"),
|
||
patch.object(client, "_fetch_spec", return_value=_SPEC),
|
||
):
|
||
with self.assertRaisesRegex(FactoryMesError, "接口定义之外"):
|
||
client.call(
|
||
"qm_ftestwork_read",
|
||
arguments={"batch": "B1", "unexpected": "x"},
|
||
)
|
||
|
||
|
||
class ExternalSystemToolSafetyTests(unittest.TestCase):
|
||
def test_status_tool_remains_registered_without_callable_connection(self):
|
||
from core.tool_registry import ToolContext, build_tools
|
||
|
||
uid = uuid.uuid4()
|
||
task_id = uuid.uuid4()
|
||
with (
|
||
tempfile.TemporaryDirectory() as tmp,
|
||
patch.dict(
|
||
os.environ,
|
||
{"DOCUMENT_SEARCH_API_KEY": "", "MP_API_KEY": ""},
|
||
clear=False,
|
||
),
|
||
patch(
|
||
"core.tool_registry._external_system_status_available",
|
||
return_value=True,
|
||
),
|
||
patch(
|
||
"core.tool_registry._external_systems_available",
|
||
return_value=False,
|
||
),
|
||
patch("core.tool_registry.smtp_configured", return_value=False),
|
||
patch("core.tool_registry.wechat_push_available", return_value=False),
|
||
patch("core.tool_registry.lfasr_configured", return_value=False),
|
||
patch("core.tool_registry.BochaConfig.load", return_value=None),
|
||
):
|
||
root = Path(tmp)
|
||
tools = build_tools(
|
||
ToolContext(
|
||
tool_base=root,
|
||
ur_path=root,
|
||
working_dir_path=root,
|
||
task_id=task_id,
|
||
uid=uid,
|
||
cfg={},
|
||
caps=SimpleNamespace(enable_run_python=False),
|
||
skills=SimpleNamespace(skills={}),
|
||
cancel_check=None,
|
||
scheduled_run=True,
|
||
deferred_actions=SimpleNamespace(),
|
||
ark_cfg=None,
|
||
img_provider="",
|
||
img_key="",
|
||
img_cfg=None,
|
||
img_provider_cfg=None,
|
||
video_variant="",
|
||
office_to_pdf_available=False,
|
||
)
|
||
)
|
||
self.assertIn("external_system_list", tools)
|
||
self.assertNotIn("external_system_search", tools)
|
||
self.assertNotIn("external_system_call", tools)
|
||
|
||
def test_tools_are_scoped_to_constructor_user(self):
|
||
from tools.external_systems import ExternalSystemListTool
|
||
|
||
uid = uuid.uuid4()
|
||
with patch(
|
||
"tools.external_systems.list_external_systems",
|
||
return_value=[
|
||
{
|
||
"external_system_id": str(uuid.uuid4()),
|
||
"status": "active",
|
||
"username_masked": "me***r",
|
||
},
|
||
{
|
||
"external_system_id": str(uuid.uuid4()),
|
||
"status": "needs_reverify",
|
||
"username_masked": "me***r",
|
||
},
|
||
],
|
||
) as listed:
|
||
output = ExternalSystemListTool(uid).execute()
|
||
listed.assert_called_once_with(uid)
|
||
self.assertNotIn("password", output.lower())
|
||
payload = json.loads(output)
|
||
self.assertEqual(len(payload["systems"]), 2)
|
||
self.assertIn("测试连接", payload["systems"][1]["status_guidance"])
|
||
|
||
def test_search_returns_admin_guidance_with_recommended_operations(self):
|
||
from tools.external_systems import ExternalSystemSearchTool
|
||
|
||
uid = uuid.uuid4()
|
||
client = SimpleNamespace(
|
||
cfg=SimpleNamespace(
|
||
query_guidance="统计查询先查看数据集目录",
|
||
recommended_operation_ids=("bi_dataset_list",),
|
||
),
|
||
search=lambda query, limit: [{"operation_id": "bi_dataset_list"}],
|
||
)
|
||
with patch(
|
||
"tools.external_systems._row_and_client",
|
||
return_value=(SimpleNamespace(), client),
|
||
):
|
||
output = ExternalSystemSearchTool(uid).execute(str(uuid.uuid4()), "产量")
|
||
payload = json.loads(output)
|
||
self.assertEqual(payload["query_guidance"], "统计查询先查看数据集目录")
|
||
self.assertEqual(payload["recommended_operation_ids"], ["bi_dataset_list"])
|
||
|
||
def test_large_call_spills_and_result_reader_pages_without_data_loss(self):
|
||
from tools.external_systems import (
|
||
ExternalSystemCallTool,
|
||
ExternalSystemResultExportTool,
|
||
ExternalSystemResultReadTool,
|
||
)
|
||
|
||
uid = uuid.uuid4()
|
||
task_id = uuid.uuid4()
|
||
system_id = str(uuid.uuid4())
|
||
budget: dict[str, int] = {}
|
||
client = SimpleNamespace(
|
||
cfg=SimpleNamespace(max_result_bytes=512, max_total_result_bytes=4096),
|
||
call=lambda *args, **kwargs: {
|
||
"operation_id": "detail_list",
|
||
"status_code": 200,
|
||
"truncated": False,
|
||
"data": [{"id": index, "noise": "x" * 100} for index in range(20)],
|
||
},
|
||
)
|
||
with (
|
||
tempfile.TemporaryDirectory() as tmp,
|
||
patch(
|
||
"tools.external_systems._row_and_client",
|
||
return_value=(SimpleNamespace(), client),
|
||
),
|
||
):
|
||
call_tool = ExternalSystemCallTool(
|
||
uid,
|
||
task_id=task_id,
|
||
result_budget=budget,
|
||
base_dir=Path(tmp),
|
||
)
|
||
spilled = json.loads(call_tool.execute(system_id, "detail_list"))
|
||
self.assertFalse(spilled["inline_complete"])
|
||
self.assertFalse(spilled["truncated"])
|
||
self.assertRegex(spilled["result_ref"], r"^extres_[0-9a-f]{32}$")
|
||
self.assertNotIn("x" * 100, json.dumps(spilled))
|
||
cache_path = (
|
||
Path(tmp)
|
||
/ ".zcbot_cache"
|
||
/ str(task_id)
|
||
/ "external_results"
|
||
/ f"{spilled['result_ref']}.json"
|
||
)
|
||
self.assertTrue(cache_path.is_file())
|
||
|
||
read_tool = ExternalSystemResultReadTool(
|
||
uid,
|
||
task_id=task_id,
|
||
result_budget=budget,
|
||
base_dir=Path(tmp),
|
||
)
|
||
page = json.loads(
|
||
read_tool.execute(
|
||
spilled["result_ref"],
|
||
json_pointer="/data",
|
||
offset=5,
|
||
limit=2,
|
||
fields=["id"],
|
||
)
|
||
)
|
||
other_task = ExternalSystemResultReadTool(
|
||
uid,
|
||
task_id=uuid.uuid4(),
|
||
base_dir=Path(tmp),
|
||
)
|
||
cross_task = other_task.execute(spilled["result_ref"])
|
||
|
||
export_tool = ExternalSystemResultExportTool(
|
||
uid,
|
||
task_id=task_id,
|
||
base_dir=Path(tmp),
|
||
user_root=Path(tmp),
|
||
)
|
||
exported = export_tool.execute(
|
||
spilled["result_ref"],
|
||
filename="detail_snapshot.json",
|
||
)
|
||
export_path = Path(tmp) / "data" / "external" / "detail_snapshot.json"
|
||
export_payload = json.loads(export_path.read_text(encoding="utf-8"))
|
||
self.assertEqual(page["total_items"], 20)
|
||
self.assertEqual(page["data"], [{"id": 5}, {"id": 6}])
|
||
self.assertTrue(page["has_more"])
|
||
self.assertIn("不存在或已过期", cross_task)
|
||
self.assertEqual(
|
||
exported.artifacts[0].path, "data/external/detail_snapshot.json"
|
||
)
|
||
self.assertEqual(export_payload["_zcbot"]["result_ref"], spilled["result_ref"])
|
||
self.assertEqual(
|
||
export_payload["_zcbot"]["provenance"]["operation_id"],
|
||
"detail_list",
|
||
)
|
||
self.assertEqual(len(export_payload["response"]["data"]), 20)
|
||
|
||
def test_result_store_reads_legacy_cache_location(self):
|
||
from core.external_systems.results import ExternalResultStore
|
||
|
||
with tempfile.TemporaryDirectory() as tmp:
|
||
task_id = str(uuid.uuid4())
|
||
store = ExternalResultStore(Path(tmp), task_id)
|
||
stored = store.store("system-1", {"data": [1, 2, 3]})
|
||
current = store.root / f"{stored.result_ref}.json"
|
||
legacy = store.legacy_root / current.name
|
||
legacy.parent.mkdir(parents=True)
|
||
current.replace(legacy)
|
||
|
||
loaded = store.load(stored.result_ref)
|
||
|
||
self.assertEqual(loaded["result"], {"data": [1, 2, 3]})
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|