creator-hub/browser_gateway/test_agent.py

"""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

import websocket
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"ws://127.0.0.1:{self.port}/v1/agent"

    def _process_request(self, connection, request):
        if self.reject:
            from websockets.http11 import Response
            return Response(403, "Forbidden", 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

    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_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")
        finally:
            agent.stop()
            runner.join(timeout=5)
            platform.close()


class OutboundAgentTaskTests(unittest.TestCase):
    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()