163 lines
6.3 KiB
Python
163 lines
6.3 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import unittest
|
|
|
|
import httpx
|
|
|
|
from agent_integration.client import OpenAgentAPIClient, OpenAgentAPIError, OpenAgentEvent
|
|
|
|
|
|
class AgentClientTests(unittest.TestCase):
|
|
def test_create_session_uses_auth_and_matching_csrf_tokens(self):
|
|
seen = {}
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
seen["method"] = request.method
|
|
seen["url"] = str(request.url)
|
|
seen["authorization"] = request.headers.get("authorization")
|
|
seen["csrf"] = request.headers.get("x-csrf-token")
|
|
seen["cookie"] = request.headers.get("cookie")
|
|
seen["payload"] = json.loads(request.content.decode("utf-8"))
|
|
return httpx.Response(200, json={"session_id": "open_sess_123", "status": "active"})
|
|
|
|
with httpx.Client(transport=httpx.MockTransport(handler)) as http_client:
|
|
client = OpenAgentAPIClient(
|
|
base_url="https://superagent.nianxx.cn/",
|
|
api_key="df_open_test",
|
|
csrf_token="csrf-test-token",
|
|
http_client=http_client,
|
|
)
|
|
response = client.create_session(
|
|
external_subject_id="customer-001",
|
|
idempotency_key="session-key",
|
|
metadata={"source": "unittest"},
|
|
)
|
|
|
|
self.assertEqual(response["session_id"], "open_sess_123")
|
|
self.assertEqual(
|
|
seen,
|
|
{
|
|
"method": "POST",
|
|
"url": "https://superagent.nianxx.cn/api/open/agent-sessions",
|
|
"authorization": "Bearer df_open_test",
|
|
"csrf": "csrf-test-token",
|
|
"cookie": "csrf_token=csrf-test-token",
|
|
"payload": {
|
|
"external_subject_id": "customer-001",
|
|
"idempotency_key": "session-key",
|
|
"metadata": {"source": "unittest"},
|
|
},
|
|
},
|
|
)
|
|
|
|
def test_stream_parser_preserves_event_id_retry_and_multiline_json(self):
|
|
seen = {}
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
self.assertEqual(request.headers.get("accept"), "text/event-stream")
|
|
seen["include_trace"] = request.url.params.get("include_trace")
|
|
return httpx.Response(
|
|
200,
|
|
headers={"content-type": "text/event-stream"},
|
|
content=(
|
|
": keepalive\r\n"
|
|
"id: evt-1\r\n"
|
|
"event: message.delta\r\n"
|
|
"retry: 1500\r\n"
|
|
"data: {\"content\":\r\n"
|
|
"data: \"你\"}\r\n\r\n"
|
|
"event: run.completed\r\n"
|
|
"data: {\"status\":\"completed\"}\r\n\r\n"
|
|
).encode("utf-8"),
|
|
)
|
|
|
|
with httpx.Client(transport=httpx.MockTransport(handler)) as http_client:
|
|
client = OpenAgentAPIClient(
|
|
base_url="https://superagent.nianxx.cn",
|
|
api_key="df_open_test",
|
|
http_client=http_client,
|
|
)
|
|
events = list(
|
|
client.stream_message(
|
|
"open_sess_123",
|
|
"你好",
|
|
include_trace=True,
|
|
)
|
|
)
|
|
|
|
self.assertEqual(
|
|
events,
|
|
[
|
|
OpenAgentEvent(
|
|
event="message.delta",
|
|
data={"content": "你"},
|
|
event_id="evt-1",
|
|
retry_ms=1500,
|
|
),
|
|
OpenAgentEvent(event="run.completed", data={"status": "completed"}),
|
|
],
|
|
)
|
|
self.assertEqual(seen["include_trace"], "true")
|
|
|
|
def test_non_stream_routes_support_x_api_key_auth(self):
|
|
calls = []
|
|
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
calls.append((request.method, request.url.path, request.headers.get("x-deerflow-open-api-key")))
|
|
if request.url.path.endswith("/messages"):
|
|
return httpx.Response(200, json={"run_id": "run_123", "status": "running"})
|
|
if request.url.path.endswith("/cancel"):
|
|
return httpx.Response(200, json={"run_id": "run_123", "status": "cancelling"})
|
|
return httpx.Response(200, json={"run_id": "run_123", "status": "completed"})
|
|
|
|
with httpx.Client(transport=httpx.MockTransport(handler)) as http_client:
|
|
client = OpenAgentAPIClient(
|
|
base_url="https://superagent.nianxx.cn",
|
|
api_key="df_open_test",
|
|
auth_mode="x-api-key",
|
|
http_client=http_client,
|
|
)
|
|
client.send_message("open_sess_123", "hello")
|
|
client.get_run("open_sess_123", "run_123")
|
|
client.cancel_run("open_sess_123", "run_123")
|
|
|
|
self.assertEqual(
|
|
calls,
|
|
[
|
|
("POST", "/api/open/agent-sessions/open_sess_123/messages", "df_open_test"),
|
|
("GET", "/api/open/agent-sessions/open_sess_123/runs/run_123", "df_open_test"),
|
|
("POST", "/api/open/agent-sessions/open_sess_123/runs/run_123/cancel", "df_open_test"),
|
|
],
|
|
)
|
|
|
|
def test_api_error_is_structured_without_exposing_api_key(self):
|
|
def handler(request: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(
|
|
409,
|
|
headers={"x-request-id": "request-123"},
|
|
json={"detail": "Active run"},
|
|
)
|
|
|
|
secret = "df_open_secret_value"
|
|
with httpx.Client(transport=httpx.MockTransport(handler)) as http_client:
|
|
client = OpenAgentAPIClient(
|
|
base_url="https://superagent.nianxx.cn",
|
|
api_key=secret,
|
|
http_client=http_client,
|
|
)
|
|
with self.assertRaises(OpenAgentAPIError) as caught:
|
|
client.send_message("open_sess_123", "hello")
|
|
|
|
error = caught.exception
|
|
self.assertEqual(error.status_code, 409)
|
|
self.assertEqual(error.detail, "Active run")
|
|
self.assertEqual(error.request_id, "request-123")
|
|
self.assertTrue(error.active_run_conflict)
|
|
self.assertFalse(error.retryable)
|
|
self.assertNotIn(secret, str(error))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main(verbosity=2)
|