Add worker gateway and websocket transport flow

2026-07-03 21:30:17 +03:00
parent b681a90eae
commit 8423d8e8ad
8 changed files with 727 additions and 12 deletions
@@ -69,3 +69,72 @@ def test_worker_registration_is_exposed_via_list_endpoint() -> None:
    assert registration.status_code == 200
    assert listing.status_code == 200
    assert any(worker["worker_id"] == "worker_home_pc" for worker in listing.json()["workers"])
+
+
+def test_worker_command_dispatch_and_poll_cycle() -> None:
+    client = TestClient(app)
+
+    registration = client.post(
+        "/workers/register",
+        json={
+            "worker_id": "worker_home_pc",
+            "name": "Home PC",
+            "capabilities": ["file.read", "command.run"],
+            "version": "0.1.0",
+            "machine": "DESKTOP-1",
+            "os": "windows",
+        },
+    ).json()
+    dispatch = client.post(
+        f"/workers/{registration['session_id']}/commands",
+        json={
+            "task_id": "task_1",
+            "node_id": "node_1",
+            "command_name": "file.read",
+            "args": {"path": "D:/test.txt"},
+            "timeout_ms": 30000,
+            "policy_context": {"approved": True},
+        },
+    )
+    polled = client.get(f"/workers/{registration['session_id']}/commands")
+
+    assert dispatch.status_code == 200
+    assert dispatch.json()["status"] == "queued"
+    assert len(polled.json()["commands"]) == 1
+
+
+def test_worker_websocket_heartbeat_and_poll() -> None:
+    client = TestClient(app)
+    registration = client.post(
+        "/workers/register",
+        json={
+            "worker_id": "worker_socket_pc",
+            "name": "Socket PC",
+            "capabilities": ["file.read"],
+            "version": "0.1.0",
+            "machine": "DESKTOP-2",
+            "os": "windows",
+        },
+    ).json()
+    client.post(
+        f"/workers/{registration['session_id']}/commands",
+        json={
+            "task_id": "task_2",
+            "node_id": "node_2",
+            "command_name": "file.read",
+            "args": {"path": "D:/socket.txt"},
+            "timeout_ms": 30000,
+            "policy_context": {"approved": True},
+        },
+    )
+
+    with client.websocket_connect(f"/workers/ws/{registration['session_id']}") as websocket:
+        websocket.send_json({"event_type": "heartbeat"})
+        heartbeat_ack = websocket.receive_json()
+        websocket.send_json({"event_type": "poll"})
+        commands = websocket.receive_json()
+
+    assert heartbeat_ack["type"] == "ack"
+    assert heartbeat_ack["status"] in {"online", "busy"}
+    assert commands["type"] == "commands"
+    assert len(commands["commands"]) == 1