273 lines
8.2 KiB
Python
273 lines
8.2 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
|
|
assert any(card["type"] == "plan_card" for card in payload["cards"])
|
|
assert any(card["type"] == "progress_card" for card in payload["cards"])
|
|
|
|
|
|
def test_chat_response_contains_ui_cards() -> None:
|
|
client = TestClient(app)
|
|
|
|
response = client.post(
|
|
"/chat",
|
|
json={
|
|
"project_id": "default",
|
|
"conversation_id": "conv_demo",
|
|
"message": "Build plan",
|
|
"attachments": [],
|
|
"mode": "auto",
|
|
"preferences": {"show_plan": True},
|
|
},
|
|
)
|
|
|
|
payload = response.json()
|
|
|
|
assert response.status_code == 200
|
|
assert payload["response_type"] == "task_started"
|
|
assert any(card["type"] == "plan_card" for card in payload["cards"])
|
|
assert any(card["type"] == "progress_card" for card in payload["cards"])
|
|
|
|
|
|
def test_task_creation_with_policy_action_returns_confirmation_card() -> None:
|
|
client = TestClient(app)
|
|
|
|
response = client.post(
|
|
"/tasks",
|
|
json={
|
|
"project_id": "default",
|
|
"goal": "Need approval",
|
|
"inputs": {
|
|
"policy_action": {
|
|
"resource": "filesystem",
|
|
"risk_level": "write",
|
|
"action_name": "file.write",
|
|
"preview_available": True,
|
|
"metadata": {"path": "D:/Projects/test.txt"},
|
|
}
|
|
},
|
|
"execution_mode": "agent_graph",
|
|
},
|
|
)
|
|
|
|
payload = response.json()
|
|
|
|
assert response.status_code == 200
|
|
assert any(card["type"] == "confirmation_card" for card in payload["cards"])
|
|
|
|
|
|
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"
|
|
assert any(card["type"] == "error_card" for card in response.json()["cards"])
|
|
|
|
|
|
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
|
|
|
|
|
|
def test_task_get_returns_artifact_card_from_worker_result() -> None:
|
|
client = TestClient(app)
|
|
created = client.post(
|
|
"/tasks",
|
|
json={
|
|
"project_id": "default",
|
|
"goal": "Artifact check",
|
|
"inputs": {},
|
|
"execution_mode": "agent_graph",
|
|
},
|
|
).json()
|
|
registration = client.post(
|
|
"/workers/register",
|
|
json={
|
|
"worker_id": "worker_artifact_pc",
|
|
"name": "Artifact PC",
|
|
"capabilities": ["file.read"],
|
|
"version": "0.1.0",
|
|
"machine": "DESKTOP-3",
|
|
"os": "windows",
|
|
},
|
|
).json()
|
|
client.post(
|
|
f"/workers/{registration['session_id']}/commands",
|
|
json={
|
|
"task_id": created["task_id"],
|
|
"node_id": "node_artifact",
|
|
"command_name": "file.read",
|
|
"args": {"path": "D:/artifact.txt"},
|
|
},
|
|
)
|
|
with client.websocket_connect(f"/workers/ws/{registration['session_id']}") as websocket:
|
|
websocket.send_json({"event_type": "poll"})
|
|
commands = websocket.receive_json()["commands"]
|
|
websocket.send_json(
|
|
{
|
|
"event_type": "result",
|
|
"command_id": commands[0]["command_id"],
|
|
"task_id": created["task_id"],
|
|
"tool": "file.read",
|
|
"status": "success",
|
|
"started_at": "2026-07-03T10:00:00+00:00",
|
|
"finished_at": "2026-07-03T10:00:01+00:00",
|
|
"duration_ms": 1000,
|
|
"stdout": "",
|
|
"stderr": "",
|
|
"result": {"content": "ok"},
|
|
"artifacts": [{"type": "file", "path": "D:/artifact.txt"}],
|
|
"error": None,
|
|
}
|
|
)
|
|
websocket.receive_json()
|
|
|
|
response = client.get(f"/tasks/{created['task_id']}")
|
|
|
|
assert response.status_code == 200
|
|
assert any(card["type"] == "artifact_card" for card in response.json()["cards"])
|
|
|
|
|
|
def test_task_events_stream_returns_normalized_event_payload() -> None:
|
|
client = TestClient(app)
|
|
created = client.post(
|
|
"/tasks",
|
|
json={
|
|
"project_id": "default",
|
|
"goal": "Event check",
|
|
"inputs": {},
|
|
"execution_mode": "agent_graph",
|
|
},
|
|
).json()
|
|
|
|
with client.stream("GET", f"/tasks/{created['task_id']}/events") as response:
|
|
body = response.read().decode()
|
|
|
|
assert response.status_code == 200
|
|
assert "event_type" in body
|
|
assert "progress_card_ref" in body or "task_planned" in body
|