forked from xiaohei/taiji-AI-PAD
313 lines
13 KiB
Python
313 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""End-to-end flow tester for taiji-AI-PAD services."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import sys
|
|
import uuid
|
|
from dataclasses import dataclass
|
|
from typing import Any, Dict, Iterable, Optional
|
|
|
|
import requests
|
|
|
|
|
|
@dataclass
|
|
class ServiceConfig:
|
|
data_ingestion_url: str = "http://localhost:8001"
|
|
mcp_server_url: str = "http://localhost:8002"
|
|
timeout: int = 30
|
|
|
|
|
|
class ApiFlowTester:
|
|
def __init__(self, config: ServiceConfig, verbose: bool = True) -> None:
|
|
self.config = config
|
|
self.session = requests.Session()
|
|
self.verbose = verbose
|
|
|
|
def run(self, skip_data_ingestion: bool, skip_mcp: bool) -> None:
|
|
if not skip_data_ingestion:
|
|
self._log("Running Data Ingestion flow")
|
|
self._test_data_ingestion_flow()
|
|
else:
|
|
self._log("Skipping Data Ingestion flow")
|
|
|
|
if not skip_mcp:
|
|
self._log("Running MCP Server flow")
|
|
self._test_mcp_flow()
|
|
else:
|
|
self._log("Skipping MCP Server flow")
|
|
|
|
def _test_data_ingestion_flow(self) -> None:
|
|
base = self.config.data_ingestion_url
|
|
|
|
self._log("Checking Data Ingestion health endpoint")
|
|
health = self._json_request("GET", f"{base}/health")
|
|
self._ensure_service_health(
|
|
health,
|
|
critical_keys=["data_ingestion"],
|
|
context="Data Ingestion",
|
|
)
|
|
|
|
self._log("Triggering RapidAPI sync job")
|
|
sync_response = self._json_request(
|
|
"POST",
|
|
f"{base}/rapidapi/sync",
|
|
params={"category": "weather", "limit": 1},
|
|
)
|
|
self._require("message" in sync_response, "RapidAPI sync did not return confirmation", sync_response)
|
|
|
|
self._log("Parsing reference OpenAPI specification")
|
|
openapi_url = "https://raw.githubusercontent.com/OAI/OpenAPI-Specification/main/examples/v3.0/petstore.yaml"
|
|
openapi = self._json_request("POST", f"{base}/openapi/parse", params={"url": openapi_url})
|
|
self._require(openapi.get("parsed_data"), "OpenAPI parse missing parsed_data", openapi)
|
|
|
|
self._log("Running APILLAMA processing step")
|
|
apillama_payload = {
|
|
"api_doc": {
|
|
"title": "Weather API",
|
|
"description": "Returns forecast information",
|
|
"parameters": [
|
|
{"name": "location", "type": "string", "description": "City name", "required": True}
|
|
],
|
|
},
|
|
"context": {"service": "weather", "version": "1.0"},
|
|
"output_format": "json_schema",
|
|
"include_examples": True,
|
|
"enhance_descriptions": True,
|
|
"validate_schema": True,
|
|
}
|
|
apillama = self._json_request("POST", f"{base}/apillama/process", json=apillama_payload)
|
|
self._require(apillama.get("processed"), "APILLAMA processing failed", apillama)
|
|
|
|
self._log("Requesting tool generation task")
|
|
tool_payload = {
|
|
"url": "https://api.example.com/weather",
|
|
"method": "GET",
|
|
"name": f"diag_get_weather_{uuid.uuid4().hex[:8]}",
|
|
"description": "Diagnostic weather fetch tool",
|
|
"parameters": [
|
|
{
|
|
"name": "location",
|
|
"type": "string",
|
|
"location": "query",
|
|
"description": "City name",
|
|
"required": True,
|
|
},
|
|
{
|
|
"name": "unit",
|
|
"type": "string",
|
|
"location": "query",
|
|
"description": "Measurement unit",
|
|
"required": False,
|
|
},
|
|
],
|
|
"request_body": None,
|
|
"responses": {
|
|
"200": {
|
|
"description": "Success",
|
|
"content": {
|
|
"application/json": {
|
|
"schema": {
|
|
"type": "object",
|
|
"properties": {
|
|
"temperature": {"type": "number"},
|
|
"condition": {"type": "string"},
|
|
},
|
|
}
|
|
}
|
|
},
|
|
}
|
|
},
|
|
"security": [],
|
|
"tags": ["weather"],
|
|
"deprecated": False,
|
|
"headers": {"Authorization": "Bearer demo-token"},
|
|
}
|
|
tool_gen = self._json_request("POST", f"{base}/tools/generate", json=tool_payload)
|
|
self._require(
|
|
tool_gen.get("message"),
|
|
"Tool generation endpoint did not acknowledge request",
|
|
tool_gen,
|
|
)
|
|
|
|
self._log("Listing generated tools")
|
|
tools = self._json_request("GET", f"{base}/tools", params={"limit": 5})
|
|
self._require(isinstance(tools, list), "Tools endpoint did not return a list", tools)
|
|
if tools:
|
|
self._log("Fetching first tool definition for verification")
|
|
first_tool = tools[0]["name"]
|
|
definition = self._json_request("GET", f"{base}/tools/{first_tool}")
|
|
self._require(
|
|
definition.get("name") == first_tool,
|
|
"Fetched tool definition does not match",
|
|
definition,
|
|
)
|
|
|
|
self._log("Reading Data Ingestion metrics endpoint")
|
|
metrics_text = self._request("GET", f"{base}/metrics").text
|
|
self._require("http_requests_total" in metrics_text, "Metrics output missing expected counters")
|
|
|
|
def _test_mcp_flow(self) -> None:
|
|
base = self.config.mcp_server_url
|
|
|
|
self._log("Checking MCP health endpoint")
|
|
health = self._json_request("GET", f"{base}/health")
|
|
self._ensure_service_health(
|
|
health,
|
|
critical_keys=["database"],
|
|
context="MCP",
|
|
)
|
|
|
|
self._log("Fetching current MCP agent list")
|
|
existing_agents = self._json_request("GET", f"{base}/agents", params={"skip": 0, "limit": 100})
|
|
self._require(isinstance(existing_agents, list), "Agent list did not return a list", existing_agents)
|
|
existing_ids = {agent.get("id") for agent in existing_agents if agent.get("id")}
|
|
|
|
self._log("Registering diagnostic agent")
|
|
agent_name = f"auto-agent-{uuid.uuid4().hex[:8]}"
|
|
agent_payload = {
|
|
"name": agent_name,
|
|
"description": "Auto-generated diagnostic agent",
|
|
"role": "assistant",
|
|
"goal": "Validate MCP server flows",
|
|
"tools": ["math_add"],
|
|
"config": {"default_model": "gpt-4o-mini"},
|
|
"capabilities": ["diagnostics"],
|
|
}
|
|
created_agent = self._json_request("POST", f"{base}/agents", json=agent_payload)
|
|
agent_id = created_agent.get("id")
|
|
self._require(agent_id, "Agent creation response missing id", created_agent)
|
|
|
|
self._log("Verifying new agent presence in list")
|
|
refreshed_agents = self._json_request("GET", f"{base}/agents", params={"skip": 0, "limit": 100})
|
|
after_ids = {agent.get("id") for agent in refreshed_agents if agent.get("id")}
|
|
self._require(agent_id in after_ids, "New agent not found in list after creation")
|
|
|
|
self._log("Fetching newly created agent details")
|
|
fetched_agent = self._json_request("GET", f"{base}/agents/{agent_id}")
|
|
self._require(fetched_agent.get("name") == agent_name, "Fetched agent does not match created agent", fetched_agent)
|
|
|
|
self._log("Executing math_add tool via MCP agent")
|
|
execution_payload = {
|
|
"jsonrpc": "2.0",
|
|
"id": f"exec-{uuid.uuid4().hex[:8]}",
|
|
"method": "tools/call",
|
|
"params": {
|
|
"tool": {"name": "math_add", "function_name": "math_add"},
|
|
"arguments": {"a": 1, "b": 2},
|
|
"context": {"session_id": f"session-{agent_id[:8]}"},
|
|
},
|
|
}
|
|
execution = self._json_request("POST", f"{base}/agents/{agent_id}/execute", json=execution_payload)
|
|
self._require(execution.get("success"), "Agent tool execution failed", execution)
|
|
self._require(execution.get("result") == 3, "math_add result mismatch", execution)
|
|
|
|
self._log("Listing MCP tools for visibility")
|
|
tools = self._json_request("GET", f"{base}/tools", params={"limit": 5})
|
|
self._require(isinstance(tools, list), "MCP tools endpoint did not return a list", tools)
|
|
|
|
self._log("Reading MCP metrics endpoint")
|
|
metrics_text = self._request("GET", f"{base}/metrics").text
|
|
self._require("http_requests_total" in metrics_text, "MCP metrics missing expected counters")
|
|
|
|
self._log("Fetching monitoring metrics snapshot")
|
|
monitoring = self._json_request("GET", f"{base}/api/v1/monitoring/metrics")
|
|
self._require(
|
|
monitoring.get("system"),
|
|
"Monitoring metrics missing system section",
|
|
monitoring,
|
|
)
|
|
|
|
def _request(
|
|
self,
|
|
method: str,
|
|
url: str,
|
|
*,
|
|
expected_status: Optional[Iterable[int]] = None,
|
|
**kwargs: Any,
|
|
) -> requests.Response:
|
|
response = self.session.request(method, url, timeout=self.config.timeout, **kwargs)
|
|
acceptable = list(expected_status) if expected_status is not None else []
|
|
if expected_status is None and not 200 <= response.status_code < 300:
|
|
raise AssertionError(
|
|
f"Request to {url} failed with status {response.status_code}: {response.text[:200]}"
|
|
)
|
|
if expected_status is not None and response.status_code not in acceptable:
|
|
raise AssertionError(
|
|
f"Request to {url} expected {acceptable} but received {response.status_code}: {response.text[:200]}"
|
|
)
|
|
return response
|
|
|
|
def _json_request(self, method: str, url: str, **kwargs: Any) -> Dict[str, Any] | Any:
|
|
response = self._request(method, url, **kwargs)
|
|
try:
|
|
return response.json()
|
|
except ValueError as exc: # pragma: no cover - defensive guard
|
|
raise AssertionError(f"Response from {url} is not valid JSON: {response.text[:200]}") from exc
|
|
|
|
def _require(self, condition: bool, message: str, payload: Optional[Any] = None) -> None:
|
|
if not condition:
|
|
detail = f" | payload={payload}" if payload is not None else ""
|
|
raise AssertionError(f"{message}{detail}")
|
|
|
|
def _ensure_service_health(
|
|
self,
|
|
payload: Dict[str, Any],
|
|
*,
|
|
critical_keys: Optional[Iterable[str]] = None,
|
|
context: str,
|
|
) -> None:
|
|
status = str(payload.get("status", "")).lower()
|
|
acceptable = {"healthy", "ok", "degraded"}
|
|
services = payload.get("services") or {}
|
|
if status not in acceptable:
|
|
raise AssertionError(f"{context} health status unacceptable: {status} | payload={payload}")
|
|
if critical_keys:
|
|
missing = [svc for svc in critical_keys if services.get(svc) not in {"healthy", "ok"}]
|
|
if missing:
|
|
raise AssertionError(
|
|
f"{context} critical services unhealthy: {missing} | payload={payload}"
|
|
)
|
|
degraded = [name for name, svc_status in services.items() if svc_status == "degraded"]
|
|
if degraded:
|
|
self._log(
|
|
f"{context} warning: degraded dependencies detected: {', '.join(degraded)}"
|
|
)
|
|
|
|
def _log(self, message: str) -> None:
|
|
if self.verbose:
|
|
print(f"[api-flow] {message}")
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Validate taiji-AI-PAD API flows")
|
|
parser.add_argument("--data-ingestion-url", default="http://localhost:8001", help="Data Ingestion base URL")
|
|
parser.add_argument("--mcp-server-url", default="http://localhost:8002", help="MCP Server base URL")
|
|
parser.add_argument("--timeout", type=int, default=30, help="HTTP timeout in seconds")
|
|
parser.add_argument("--skip-data", action="store_true", help="Skip Data Ingestion flow")
|
|
parser.add_argument("--skip-mcp", action="store_true", help="Skip MCP flow")
|
|
parser.add_argument("--quiet", action="store_true", help="Suppress verbose logs")
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> None:
|
|
args = parse_args()
|
|
config = ServiceConfig(
|
|
data_ingestion_url=args.data_ingestion_url,
|
|
mcp_server_url=args.mcp_server_url,
|
|
timeout=args.timeout,
|
|
)
|
|
tester = ApiFlowTester(config, verbose=not args.quiet)
|
|
tester.run(skip_data_ingestion=args.skip_data, skip_mcp=args.skip_mcp)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
main()
|
|
except (AssertionError, requests.RequestException) as exc:
|
|
print(f"[api-flow] ❌ {exc}", file=sys.stderr)
|
|
sys.exit(1)
|
|
except KeyboardInterrupt:
|
|
print("[api-flow] Interrupted", file=sys.stderr)
|
|
sys.exit(130)
|