141 lines
4.0 KiB
Python
141 lines
4.0 KiB
Python
from fastapi.testclient import TestClient
|
|
|
|
from ai_orchestrator.main import app
|
|
|
|
|
|
def test_health_endpoint_returns_ok() -> None:
|
|
client = TestClient(app)
|
|
|
|
response = client.get("/health")
|
|
|
|
assert response.status_code == 200
|
|
assert response.json() == {"status": "ok"}
|
|
|
|
|
|
def test_post_tasks_creates_planned_task() -> None:
|
|
client = TestClient(app)
|
|
|
|
response = client.post(
|
|
"/tasks",
|
|
json={
|
|
"project_id": "default",
|
|
"goal": "Prepare a plan",
|
|
"inputs": {"source": "test"},
|
|
"execution_mode": "agent_graph",
|
|
},
|
|
)
|
|
|
|
payload = response.json()
|
|
|
|
assert response.status_code == 200
|
|
assert payload["status"] == "running"
|
|
assert payload["progress"]["total"] == 2
|
|
|
|
|
|
def test_cancel_endpoint_marks_task_cancelled() -> None:
|
|
client = TestClient(app)
|
|
|
|
created = client.post(
|
|
"/tasks",
|
|
json={
|
|
"project_id": "default",
|
|
"goal": "Cancel me",
|
|
"inputs": {},
|
|
"execution_mode": "agent_graph",
|
|
},
|
|
).json()
|
|
response = client.post(f"/tasks/{created['task_id']}/cancel", json={"reason": "stop"})
|
|
|
|
assert response.status_code == 200
|
|
assert response.json()["status"] == "cancelled"
|
|
|
|
|
|
def test_worker_registration_is_exposed_via_list_endpoint() -> 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",
|
|
},
|
|
)
|
|
listing = client.get("/workers")
|
|
|
|
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
|