463 lines
20 KiB
Python
463 lines
20 KiB
Python
"""Outbound agent contract tests: connection, heartbeat, task dispatch, event push."""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import threading
|
||
import time
|
||
import unittest
|
||
from importlib import import_module
|
||
from unittest.mock import Mock, patch
|
||
|
||
import websocket
|
||
from websockets.datastructures import Headers
|
||
from websockets.http11 import Response
|
||
from websockets.sync.server import serve
|
||
|
||
agent_module = import_module(f"{__package__}.agent")
|
||
gateway_module = import_module(f"{__package__}.server.http")
|
||
|
||
AgentConfig = agent_module.AgentConfig
|
||
OutboundAgent = agent_module.OutboundAgent
|
||
RequestError = gateway_module.RequestError
|
||
|
||
|
||
class FakePlatform:
|
||
"""Minimal platform-side WS peer capturing frames and replaying scripted ones."""
|
||
|
||
def __init__(self, *, reject: bool = False) -> None:
|
||
self.reject = reject
|
||
self.server = serve(self._handler, "127.0.0.1", 0, process_request=self._process_request)
|
||
self.port = self.server.socket.getsockname()[1]
|
||
self.thread = threading.Thread(target=self.server.serve_forever, daemon=True)
|
||
self.thread.start()
|
||
self.received: list[dict] = []
|
||
self.lock = threading.Lock()
|
||
self.connections: list = []
|
||
self.on_task: dict[int, dict] = {}
|
||
self.reject_task_ids: set[int] = set()
|
||
self.keep_open = threading.Event()
|
||
|
||
def url(self) -> str:
|
||
return f"http://127.0.0.1:{self.port}"
|
||
|
||
def _process_request(self, connection, request):
|
||
if request.path != "/v1/agent":
|
||
return Response(404, "Not Found", Headers(), b"invalid agent path\r\n")
|
||
if self.reject:
|
||
return Response(403, "Forbidden", Headers(), b"invalid access key\r\n")
|
||
return None
|
||
|
||
def _handler(self, ws) -> None:
|
||
with self.lock:
|
||
self.connections.append(ws)
|
||
try:
|
||
while True:
|
||
raw = ws.recv(timeout=5)
|
||
if not isinstance(raw, str):
|
||
continue
|
||
frame = json.loads(raw)
|
||
with self.lock:
|
||
self.received.append(frame)
|
||
if frame.get("type") == "task":
|
||
task_id = frame["id"]
|
||
if task_id in self.reject_task_ids:
|
||
ws.send(json.dumps({"type": "result", "id": task_id, "status": 409, "error": "duplicate task id"}))
|
||
continue
|
||
scripted = self.on_task.get(task_id)
|
||
if scripted is not None:
|
||
ws.send(json.dumps({"type": "result", "id": task_id, **scripted}))
|
||
except Exception:
|
||
pass # 连接关闭/超时即结束 handler；websockets 连接线程非 daemon，必须自然退出
|
||
|
||
def wait_for(self, predicate, timeout: float = 3.0) -> dict | None:
|
||
deadline = time.monotonic() + timeout
|
||
while time.monotonic() < deadline:
|
||
with self.lock:
|
||
for frame in self.received:
|
||
if predicate(frame):
|
||
return frame
|
||
time.sleep(0.01)
|
||
return None
|
||
|
||
def latest_connection(self):
|
||
with self.lock:
|
||
return self.connections[-1] if self.connections else None
|
||
|
||
def close(self) -> None:
|
||
self.keep_open.set()
|
||
self.server.shutdown()
|
||
for ws in list(self.connections):
|
||
try:
|
||
ws.close()
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
def make_gateway() -> Mock:
|
||
gateway = Mock()
|
||
gateway.token = "unused"
|
||
manager = Mock()
|
||
# 事件桥依赖真实 Condition/版本语义；替身必须与 SubscriptionManager 同构。
|
||
manager.changed = threading.Condition()
|
||
manager.version = 0
|
||
manager._items = {}
|
||
gateway.subscriptions = manager
|
||
return gateway
|
||
|
||
|
||
def make_agent(gateway: Mock, platform: FakePlatform) -> OutboundAgent:
|
||
return OutboundAgent(gateway, AgentConfig(
|
||
platform_url=platform.url(),
|
||
access_key="agent-test-key-0001",
|
||
heartbeat_interval=0.1,
|
||
reconnect_base=0.05,
|
||
reconnect_max=0.2,
|
||
))
|
||
|
||
|
||
class OutboundAgentConnectionTests(unittest.TestCase):
|
||
def test_platform_root_url_resolves_to_agent_endpoint(self) -> None:
|
||
cases = (
|
||
("http://127.0.0.1:3010", "ws://127.0.0.1:3010/v1/agent"),
|
||
("http://127.0.0.1:3010/", "ws://127.0.0.1:3010/v1/agent"),
|
||
("https://platform.example", "wss://platform.example/v1/agent"),
|
||
("https://platform.example/", "wss://platform.example/v1/agent"),
|
||
("https://platform.example/creator/", "wss://platform.example/creator/v1/agent"),
|
||
)
|
||
for root, expected in cases:
|
||
with self.subTest(root=root):
|
||
agent = OutboundAgent(make_gateway(), AgentConfig(platform_url=root, access_key="test-key"))
|
||
with patch.object(agent_module.websocket, "create_connection") as connect:
|
||
agent._connect()
|
||
self.assertEqual(connect.call_args.args[0], expected)
|
||
|
||
def test_heartbeats_flow_with_version_until_stopped(self) -> None:
|
||
platform = FakePlatform()
|
||
agent = make_agent(make_gateway(), platform)
|
||
runner = threading.Thread(target=agent.run, daemon=True)
|
||
runner.start()
|
||
try:
|
||
beat = platform.wait_for(lambda f: f.get("type") == "heartbeat")
|
||
self.assertIsNotNone(beat, "first heartbeat missing")
|
||
self.assertEqual(beat.get("version"), agent.config.version)
|
||
second = platform.wait_for(
|
||
lambda f: f.get("type") == "heartbeat" and f is not beat, timeout=3.0
|
||
)
|
||
self.assertIsNotNone(second, "periodic heartbeat missing")
|
||
finally:
|
||
agent.stop()
|
||
runner.join(timeout=5)
|
||
platform.close()
|
||
self.assertFalse(runner.is_alive())
|
||
|
||
def test_auth_failure_is_reported_and_retried_with_backoff(self) -> None:
|
||
platform = FakePlatform(reject=True)
|
||
agent = make_agent(make_gateway(), platform)
|
||
runner = threading.Thread(target=agent.run, daemon=True)
|
||
runner.start()
|
||
try:
|
||
deadline = time.monotonic() + 3.0
|
||
while time.monotonic() < deadline and agent.reconnect_attempts < 2:
|
||
time.sleep(0.01)
|
||
self.assertGreaterEqual(agent.reconnect_attempts, 2, "auth failure must trigger retry")
|
||
self.assertTrue(agent.last_error, "auth failure must be reported")
|
||
self.assertIn("403", agent.last_error, "fixture must return the intended rejection, not an internal error")
|
||
finally:
|
||
agent.stop()
|
||
runner.join(timeout=5)
|
||
platform.close()
|
||
|
||
|
||
class OutboundAgentTaskTests(unittest.TestCase):
|
||
def test_json_object_payload_is_forwarded_without_decoding(self) -> None:
|
||
gateway = make_gateway()
|
||
agent = OutboundAgent(gateway, AgentConfig(platform_url="http://127.0.0.1:3010", access_key="test-key"))
|
||
payload = {"alias": "account-a", "target_uid": "7354367890123456789", "text": "测试消息", "options": {"enabled": True}}
|
||
with patch.object(gateway_module, "route_gateway_request", return_value={"success": True}) as route:
|
||
status, body, error = agent._execute_task("POST", "/v1/social/action", payload)
|
||
self.assertEqual(status, 200)
|
||
self.assertEqual(json.loads(body), {"success": True})
|
||
self.assertEqual(error, "")
|
||
route.assert_called_once_with(gateway, "POST", "/v1/social/action", {}, payload)
|
||
|
||
def test_null_and_empty_object_payloads_are_valid(self) -> None:
|
||
for payload in (None, {}):
|
||
with self.subTest(payload=payload):
|
||
gateway = make_gateway()
|
||
agent = OutboundAgent(gateway, AgentConfig(platform_url="http://127.0.0.1:3010", access_key="test-key"))
|
||
with patch.object(gateway_module, "route_gateway_request", return_value={}) as route:
|
||
status, _, error = agent._execute_task("GET", "/v1/info", payload)
|
||
self.assertEqual(status, 200)
|
||
self.assertEqual(error, "")
|
||
route.assert_called_once_with(gateway, "GET", "/v1/info", {}, {})
|
||
|
||
def test_non_object_payloads_are_rejected_before_routing(self) -> None:
|
||
for payload in ('{"alias":"account-a"}', "", [], [1], 0, 1, False, True):
|
||
with self.subTest(payload=payload):
|
||
agent = OutboundAgent(make_gateway(), AgentConfig(platform_url="http://127.0.0.1:3010", access_key="test-key"))
|
||
with patch.object(gateway_module, "route_gateway_request", return_value={}) as route:
|
||
status, body, error = agent._execute_task("POST", "/v1/social/action", payload)
|
||
self.assertEqual(status, 400)
|
||
self.assertEqual(body, b"")
|
||
self.assertIn("JSON object", error)
|
||
route.assert_not_called()
|
||
|
||
def test_post_task_frame_preserves_json_object_and_result(self) -> None:
|
||
platform = FakePlatform()
|
||
gateway = make_gateway()
|
||
agent = make_agent(gateway, platform)
|
||
payload = {"alias": "account-a", "target_uid": "7354367890123456789", "text": "测试消息"}
|
||
result_body = {"success": True, "message_id": "7354367890123456790"}
|
||
with patch.object(gateway_module, "route_gateway_request", return_value=result_body) as route:
|
||
runner = threading.Thread(target=agent.run, daemon=True)
|
||
runner.start()
|
||
try:
|
||
self.wait_until_connected(agent)
|
||
platform.latest_connection().send(json.dumps({
|
||
"type": "task", "id": 17, "method": "POST", "path": "/v1/social/action", "payload": payload,
|
||
}, ensure_ascii=False))
|
||
result = platform.wait_for(lambda f: f.get("type") == "result" and f.get("id") == 17)
|
||
self.assertIsNotNone(result, "POST result missing")
|
||
self.assertEqual(result.get("status"), 200)
|
||
import base64
|
||
self.assertEqual(json.loads(base64.b64decode(result.get("body"))), result_body)
|
||
self.assertEqual(result.get("error"), "")
|
||
route.assert_called_once_with(gateway, "POST", "/v1/social/action", {}, payload)
|
||
finally:
|
||
agent.stop()
|
||
runner.join(timeout=5)
|
||
platform.close()
|
||
|
||
def test_task_dispatches_route_and_returns_result(self) -> None:
|
||
platform = FakePlatform()
|
||
gateway = make_gateway()
|
||
gateway.info.return_value = {"service": "browser-gateway", "node_id": "node-a"}
|
||
agent = make_agent(gateway, platform)
|
||
runner = threading.Thread(target=agent.run, daemon=True)
|
||
runner.start()
|
||
try:
|
||
self.wait_until_connected(agent)
|
||
platform.latest_connection().send(json.dumps({
|
||
"type": "task", "id": 7, "method": "GET", "path": "/v1/info", "payload": None,
|
||
}))
|
||
result = platform.wait_for(lambda f: f.get("type") == "result" and f.get("id") == 7)
|
||
self.assertIsNotNone(result, "task result missing")
|
||
self.assertEqual(result.get("status"), 200)
|
||
import base64
|
||
self.assertEqual(
|
||
json.loads(base64.b64decode(result.get("body"))),
|
||
{"service": "browser-gateway", "node_id": "node-a"},
|
||
)
|
||
finally:
|
||
agent.stop()
|
||
runner.join(timeout=5)
|
||
platform.close()
|
||
|
||
def test_unknown_route_returns_explicit_404_result(self) -> None:
|
||
platform = FakePlatform()
|
||
gateway = make_gateway()
|
||
agent = make_agent(gateway, platform)
|
||
runner = threading.Thread(target=agent.run, daemon=True)
|
||
runner.start()
|
||
try:
|
||
self.wait_until_connected(agent)
|
||
platform.latest_connection().send(json.dumps({
|
||
"type": "task", "id": 8, "method": "GET", "path": "/v1/nothing", "payload": None,
|
||
}))
|
||
result = platform.wait_for(lambda f: f.get("type") == "result" and f.get("id") == 8)
|
||
self.assertIsNotNone(result)
|
||
self.assertEqual(result.get("status"), 404)
|
||
finally:
|
||
agent.stop()
|
||
runner.join(timeout=5)
|
||
platform.close()
|
||
|
||
def test_duplicate_task_id_is_rejected_idempotently(self) -> None:
|
||
platform = FakePlatform()
|
||
gateway = make_gateway()
|
||
gateway.list_browsers.return_value = []
|
||
agent = make_agent(gateway, platform)
|
||
runner = threading.Thread(target=agent.run, daemon=True)
|
||
runner.start()
|
||
try:
|
||
self.wait_until_connected(agent)
|
||
frame = {"type": "task", "id": 9, "method": "GET", "path": "/v1/browsers", "payload": None}
|
||
platform.latest_connection().send(json.dumps(frame))
|
||
first = platform.wait_for(lambda f: f.get("type") == "result" and f.get("id") == 9)
|
||
self.assertIsNotNone(first)
|
||
platform.latest_connection().send(json.dumps(frame))
|
||
second = platform.wait_for(
|
||
lambda f: f.get("type") == "result" and f.get("id") == 9 and f is not first,
|
||
timeout=3.0,
|
||
)
|
||
self.assertIsNotNone(second)
|
||
self.assertEqual(second.get("status"), 409)
|
||
self.assertIn("duplicate", second.get("error", ""))
|
||
finally:
|
||
agent.stop()
|
||
runner.join(timeout=5)
|
||
platform.close()
|
||
|
||
def test_route_exception_maps_to_error_result(self) -> None:
|
||
platform = FakePlatform()
|
||
gateway = make_gateway()
|
||
gateway.list_browsers.side_effect = ValueError("invalid alias")
|
||
agent = make_agent(gateway, platform)
|
||
runner = threading.Thread(target=agent.run, daemon=True)
|
||
runner.start()
|
||
try:
|
||
self.wait_until_connected(agent)
|
||
platform.latest_connection().send(json.dumps({
|
||
"type": "task", "id": 10, "method": "GET", "path": "/v1/browsers", "payload": None,
|
||
}))
|
||
result = platform.wait_for(lambda f: f.get("type") == "result" and f.get("id") == 10)
|
||
self.assertIsNotNone(result)
|
||
self.assertEqual(result.get("status"), 400)
|
||
self.assertIn("invalid alias", result.get("error", ""))
|
||
finally:
|
||
agent.stop()
|
||
runner.join(timeout=5)
|
||
platform.close()
|
||
|
||
@staticmethod
|
||
def wait_until_connected(agent: OutboundAgent, timeout: float = 5.0) -> None:
|
||
deadline = time.monotonic() + timeout
|
||
while time.monotonic() < deadline and not agent.connected.is_set():
|
||
time.sleep(0.01)
|
||
if not agent.connected.is_set():
|
||
raise AssertionError("agent did not connect")
|
||
|
||
|
||
class OutboundAgentEventTests(unittest.TestCase):
|
||
def test_pending_events_are_pushed_and_ack_removes_them(self) -> None:
|
||
platform = FakePlatform()
|
||
gateway = make_gateway()
|
||
manager = gateway.subscriptions
|
||
manager.version = 1
|
||
item = Mock()
|
||
item.error = ""
|
||
item.stopped = False
|
||
item.pending.return_value = [
|
||
{"delivery_id": "d-1", "kind": "notice"},
|
||
{"delivery_id": "d-2", "kind": "checkpoint"},
|
||
]
|
||
manager._items["alias-a"] = item
|
||
manager._get.return_value = item
|
||
agent = make_agent(gateway, platform)
|
||
runner = threading.Thread(target=agent.run, daemon=True)
|
||
runner.start()
|
||
try:
|
||
OutboundAgentTaskTests.wait_until_connected(agent)
|
||
batch = platform.wait_for(lambda f: f.get("type") == "event")
|
||
self.assertIsNotNone(batch, "event batch missing")
|
||
self.assertEqual(batch["event"]["alias"], "alias-a")
|
||
self.assertEqual(len(batch["event"]["deliveries"]), 2)
|
||
platform.latest_connection().send(json.dumps({
|
||
"type": "event_ack", "alias": "alias-a", "delivery_ids": ["d-1", "d-2"],
|
||
}))
|
||
deadline = time.monotonic() + 3.0
|
||
while time.monotonic() < deadline and manager.ack.call_count == 0:
|
||
time.sleep(0.01)
|
||
manager.ack.assert_called_once()
|
||
args = manager.ack.call_args
|
||
self.assertEqual(args.args[0], "alias-a")
|
||
self.assertEqual(sorted(args.args[1]), ["d-1", "d-2"])
|
||
finally:
|
||
agent.stop()
|
||
runner.join(timeout=5)
|
||
platform.close()
|
||
|
||
def test_nack_resets_unconfirmed_batch_for_repush(self) -> None:
|
||
platform = FakePlatform()
|
||
gateway = make_gateway()
|
||
manager = gateway.subscriptions
|
||
manager.changed = threading.Condition()
|
||
item = Mock()
|
||
item.error = ""
|
||
item.stopped = False
|
||
item.pending.return_value = [{"delivery_id": "d-9", "kind": "notice"}]
|
||
manager._items["alias-a"] = item
|
||
manager._get.return_value = item
|
||
agent = make_agent(gateway, platform)
|
||
runner = threading.Thread(target=agent.run, daemon=True)
|
||
runner.start()
|
||
try:
|
||
OutboundAgentTaskTests.wait_until_connected(agent)
|
||
first = platform.wait_for(lambda f: f.get("type") == "event")
|
||
self.assertIsNotNone(first)
|
||
platform.latest_connection().send(json.dumps({
|
||
"type": "event_nack", "subscription": first["event"]["subscription"], "alias": "alias-a",
|
||
}))
|
||
repush = platform.wait_for(
|
||
lambda f: f.get("type") == "event" and f is not first, timeout=3.0
|
||
)
|
||
self.assertIsNotNone(repush, "nacked batch was not re-pushed")
|
||
self.assertEqual([d["delivery_id"] for d in repush["event"]["deliveries"]], ["d-9"])
|
||
finally:
|
||
agent.stop()
|
||
runner.join(timeout=5)
|
||
platform.close()
|
||
|
||
def test_reconnect_clears_unconfirmed_state_and_repushes(self) -> None:
|
||
platform = FakePlatform()
|
||
gateway = make_gateway()
|
||
manager = gateway.subscriptions
|
||
manager.changed = threading.Condition()
|
||
item = Mock()
|
||
item.error = ""
|
||
item.stopped = False
|
||
item.pending.return_value = [{"delivery_id": "d-7", "kind": "notice"}]
|
||
manager._items["alias-a"] = item
|
||
manager._get.return_value = item
|
||
agent = make_agent(gateway, platform)
|
||
runner = threading.Thread(target=agent.run, daemon=True)
|
||
runner.start()
|
||
try:
|
||
OutboundAgentTaskTests.wait_until_connected(agent)
|
||
first = platform.wait_for(lambda f: f.get("type") == "event")
|
||
self.assertIsNotNone(first)
|
||
# 断开连接：agent 应退避重连，并重推未确认批次
|
||
first_connection = platform.latest_connection()
|
||
first_connection.close()
|
||
deadline = time.monotonic() + 5.0
|
||
repush = None
|
||
while time.monotonic() < deadline:
|
||
repush = platform.wait_for(
|
||
lambda f: f.get("type") == "event" and f is not first, timeout=0.2
|
||
)
|
||
if repush is not None:
|
||
break
|
||
self.assertIsNotNone(repush, "unconfirmed batch not re-pushed after reconnect")
|
||
finally:
|
||
agent.stop()
|
||
runner.join(timeout=5)
|
||
platform.close()
|
||
|
||
def test_subscription_error_is_surfaced_as_error_frame(self) -> None:
|
||
platform = FakePlatform()
|
||
gateway = make_gateway()
|
||
manager = gateway.subscriptions
|
||
manager.changed = threading.Condition()
|
||
item = Mock()
|
||
item.error = "browser identity changed"
|
||
item.stopped = False
|
||
item.pending.return_value = []
|
||
manager._items["alias-a"] = item
|
||
manager._get.return_value = item
|
||
agent = make_agent(gateway, platform)
|
||
runner = threading.Thread(target=agent.run, daemon=True)
|
||
runner.start()
|
||
try:
|
||
OutboundAgentTaskTests.wait_until_connected(agent)
|
||
error = platform.wait_for(lambda f: f.get("type") == "error" and f.get("alias"))
|
||
self.assertIsNotNone(error, "subscription error frame missing")
|
||
self.assertIn("identity changed", error.get("error", ""))
|
||
finally:
|
||
agent.stop()
|
||
runner.join(timeout=5)
|
||
platform.close()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
unittest.main()
|