684 lines
27 KiB
Python
684 lines
27 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import sys
|
|
import tempfile
|
|
import unittest
|
|
import uuid
|
|
from copy import deepcopy
|
|
from pathlib import Path
|
|
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("v1:"))
|
|
self.assertNotIn("mes-password", stored)
|
|
self.assertEqual(decrypt_secret(stored), "mes-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")
|
|
|
|
|
|
def _cfg(*, allowed=frozenset(), recommended=()):
|
|
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/",
|
|
allowed_post_operations=frozenset(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),
|
|
)
|
|
|
|
|
|
_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 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"}]})
|
|
|
|
|
|
class FactoryOpenApiConnectorTests(unittest.TestCase):
|
|
def setUp(self):
|
|
from core.external_systems import factory
|
|
factory._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",
|
|
"allowed_post_operations": "bi_dataset_exec, report_preview",
|
|
"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.allowed_post_operations, {"bi_dataset_exec", "report_preview"})
|
|
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_generic_api_key_auth_uses_declared_header_without_login(self):
|
|
from core.external_systems.openapi import OpenApiClient, OpenApiConfig, _SPEC_CACHE
|
|
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 OpenApiClient, OpenApiConfig, _SPEC_CACHE
|
|
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_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_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_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_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={"wrong": "shape"},
|
|
)
|
|
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_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",
|
|
}],
|
|
) as listed:
|
|
output = ExternalSystemListTool(uid).execute()
|
|
listed.assert_called_once_with(uid)
|
|
self.assertNotIn("password", output.lower())
|
|
|
|
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()
|