File size: 8,630 Bytes
4f58e42 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 | """WebSocket integration tests.
Verifies the /ws endpoint works with correct message formats.
Auto-validators test: connect -> reset -> step -> diagnose.
Key discovery: WSResetMessage has a `data: Dict[str, Any]` field.
Task selection via WS: {"type": "reset", "data": {"task_id": "task_003"}}
"""
from __future__ import annotations
import json
import pytest
from fastapi.testclient import TestClient
from server.app import app
class TestWebSocketEndpoint:
"""Test WebSocket /ws endpoint."""
def test_ws_endpoint_exists(self) -> None:
paths = [r.path for r in app.routes if hasattr(r, "path")]
assert "/ws" in paths
def test_ws_reset_returns_observation(self) -> None:
client = TestClient(app)
with client.websocket_connect("/ws") as ws:
ws.send_json({"type": "reset"})
resp = ws.receive_json()
assert resp["type"] == "observation"
obs = resp["data"]["observation"]
assert len(obs["training_loss_history"]) == 20
assert len(obs["val_accuracy_history"]) == 20
assert len(obs["val_loss_history"]) == 20
assert obs["framework"] == "pytorch"
assert obs["epoch"] == 20
assert isinstance(obs["available_actions"], list)
assert len(obs["available_actions"]) > 0
assert obs["episode_state"]["step_count"] == 0
def test_ws_reset_with_task_selection(self) -> None:
"""Task selection via WS using data field."""
client = TestClient(app)
with client.websocket_connect("/ws") as ws:
# Task 3 is data leakage — has specific notes
ws.send_json({"type": "reset", "data": {"task_id": "task_003", "seed": 42}})
resp = ws.receive_json()
assert resp["type"] == "observation"
obs = resp["data"]["observation"]
assert "architecture upgraded" in obs.get("notes", "").lower()
assert obs["error_log"] is None # Task 3 has no error log
def test_ws_task_selection_all_tasks(self) -> None:
"""Verify all 6 tasks can be selected via WS."""
client = TestClient(app)
task_ids = ["task_001", "task_002", "task_003", "task_004", "task_005", "task_006"]
for task_id in task_ids:
with client.websocket_connect("/ws") as ws:
ws.send_json({"type": "reset", "data": {"task_id": task_id, "seed": 42}})
resp = ws.receive_json()
assert resp["type"] == "observation", f"{task_id} failed reset"
obs = resp["data"]["observation"]
assert len(obs["training_loss_history"]) == 20, f"{task_id} missing loss history"
def test_ws_step_inspect_gradients(self) -> None:
client = TestClient(app)
with client.websocket_connect("/ws") as ws:
ws.send_json({"type": "reset"})
ws.receive_json()
ws.send_json(
{"type": "step", "data": {"action_type": "inspect_gradients"}}
)
resp = ws.receive_json()
assert resp["type"] == "observation"
obs = resp["data"]["observation"]
assert len(obs["gradient_stats"]) == 4
assert obs["episode_state"]["gradients_inspected"] is True
for g in obs["gradient_stats"]:
assert "layer_name" in g
assert "mean_norm" in g
assert "is_exploding" in g
assert "is_vanishing" in g
def test_ws_full_episode_flow(self) -> None:
"""Full episode: reset -> inspect -> fix -> restart -> diagnose."""
client = TestClient(app)
with client.websocket_connect("/ws") as ws:
# Reset to task_001 (exploding gradients)
ws.send_json({"type": "reset", "data": {"task_id": "task_001", "seed": 42}})
resp = ws.receive_json()
obs = resp["data"]["observation"]
assert obs["error_log"] is not None
# Inspect gradients
ws.send_json(
{"type": "step", "data": {"action_type": "inspect_gradients"}}
)
resp = ws.receive_json()
obs = resp["data"]["observation"]
assert any(g["is_exploding"] for g in obs["gradient_stats"])
# Fix: reduce learning rate
ws.send_json(
{
"type": "step",
"data": {
"action_type": "modify_config",
"target": "learning_rate",
"value": 0.001,
},
}
)
resp = ws.receive_json()
obs = resp["data"]["observation"]
assert obs["episode_state"]["fix_action_taken"] is True
# Restart
ws.send_json({"type": "step", "data": {"action_type": "restart_run"}})
resp = ws.receive_json()
obs = resp["data"]["observation"]
assert obs["episode_state"]["restart_after_fix"] is True
# Diagnose
ws.send_json(
{
"type": "step",
"data": {
"action_type": "mark_diagnosed",
"diagnosis": "lr_too_high",
},
}
)
resp = ws.receive_json()
done = resp["data"].get("done", False)
obs = resp["data"]["observation"]
assert done or obs["episode_state"]["diagnosis_submitted"]
def test_ws_task_005_red_herrings(self) -> None:
"""Task 5 via WS — verify red herrings and correct diagnosis path."""
client = TestClient(app)
with client.websocket_connect("/ws") as ws:
ws.send_json({"type": "reset", "data": {"task_id": "task_005", "seed": 42}})
resp = ws.receive_json()
obs = resp["data"]["observation"]
# Task 5 has GPU memory warning
assert obs.get("error_log") is not None
assert obs["gpu_memory_used_gb"] > 14.0 # 91% of 16GB
# Inspect gradients — all should be non-exploding
ws.send_json(
{"type": "step", "data": {"action_type": "inspect_gradients"}}
)
resp = ws.receive_json()
obs = resp["data"]["observation"]
for g in obs["gradient_stats"]:
assert not g["is_exploding"]
# Inspect model modes — should reveal eval mode
ws.send_json(
{"type": "step", "data": {"action_type": "inspect_model_modes"}}
)
resp = ws.receive_json()
obs = resp["data"]["observation"]
assert any(v == "eval" for v in obs["model_mode_info"].values())
def test_ws_task_006_code_inspection(self) -> None:
"""Task 6 via WS — verify code inspection and fix."""
client = TestClient(app)
with client.websocket_connect("/ws") as ws:
ws.send_json({"type": "reset", "data": {"task_id": "task_006", "seed": 42}})
ws.receive_json()
# Inspect code
ws.send_json(
{"type": "step", "data": {"action_type": "inspect_code"}}
)
resp = ws.receive_json()
obs = resp["data"]["observation"]
assert obs["code_snippet"] is not None
assert obs["code_snippet"]["filename"] == "train.py"
assert obs["code_snippet"]["line_count"] > 0
def test_ws_invalid_message_returns_error(self) -> None:
client = TestClient(app)
with client.websocket_connect("/ws") as ws:
ws.send_json({"type": "reset"})
ws.receive_json()
# Wrong format — "action" instead of "data"
ws.send_json(
{"type": "step", "action": {"action_type": "inspect_gradients"}}
)
resp = ws.receive_json()
assert resp["type"] == "error"
def test_ws_step_data_batch(self) -> None:
client = TestClient(app)
with client.websocket_connect("/ws") as ws:
ws.send_json({"type": "reset"})
ws.receive_json()
ws.send_json(
{"type": "step", "data": {"action_type": "inspect_data_batch"}}
)
resp = ws.receive_json()
obs = resp["data"]["observation"]
assert obs["data_batch_stats"] is not None
assert "class_overlap_score" in obs["data_batch_stats"]
assert obs["episode_state"]["data_inspected"] is True
|