feat: prepare ARR for controlled public deployment
This commit is contained in:
Wyndham ARR
153
tests/test_agent_client.py
Normal file
153
tests/test_agent_client.py
Normal file
@@ -0,0 +1,153 @@
|
||||
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):
|
||||
def handler(request: httpx.Request) -> httpx.Response:
|
||||
self.assertEqual(request.headers.get("accept"), "text/event-stream")
|
||||
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", "你好"))
|
||||
|
||||
self.assertEqual(
|
||||
events,
|
||||
[
|
||||
OpenAgentEvent(
|
||||
event="message.delta",
|
||||
data={"content": "你"},
|
||||
event_id="evt-1",
|
||||
retry_ms=1500,
|
||||
),
|
||||
OpenAgentEvent(event="run.completed", data={"status": "completed"}),
|
||||
],
|
||||
)
|
||||
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user