File size: 11,841 Bytes
dfb775d
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
"""Coach UI deploy endpoints — env validation, spawn shim, 409 concurrency, SSE.

Mirrors the existing _SPAWN injection pattern in `tests/test_runs_sse.py`:
the three deploy spawn shims are monkeypatched per-test so we never invoke
real ssh/rsync/git/gh/httpx.
"""

from __future__ import annotations

import json
import time
from collections.abc import Iterator
from typing import Any

import pytest
from fastapi.testclient import TestClient

from mindxtrain.operator import runs as _runs
from mindxtrain.operator.app import app
from mindxtrain.operator.coach import api as coach_api

client = TestClient(app)

_GITHUB_ENV = {
    "GITHUB_TOKEN": "ghp_TEST",
    "GITHUB_REPO": "professor-codephreak/mindXtrain",
    "GITHUB_DEFAULT_BRANCH": "main",
}

_DROPLET_ENV = {
    "DROPLET_HOST": "mi300x.test",
    "DROPLET_USER": "root",
    "DROPLET_SSH_KEY": "/dev/null",
    "DROPLET_REMOTE_PATH": "/workspace/mindxtrain",
    "DROPLET_CONTAINER": "rocm/primus:v26.2",
}

_AMD_DC_ENV = {
    "AMD_DEV_CLOUD_TOKEN": "dop_v1_TEST",
    "AMD_DEV_CLOUD_SSH_KEY_ID": "56216059",
    "AMD_DEV_CLOUD_REGION": "atl1",
    "AMD_DEV_CLOUD_SIZE": "gpu-mi300x8-1536gb-devcloud",
    "AMD_DEV_CLOUD_IMAGE": "vllm-0-17-1",
}


@pytest.fixture(autouse=True)
def _restore_deploy_spawns() -> Iterator[None]:
    g = coach_api._GITHUB_SPAWN
    s = coach_api._DROPLET_SYNC_SPAWN
    p = coach_api._DROPLET_PROVISION_SPAWN
    yield
    coach_api._GITHUB_SPAWN = g
    coach_api._DROPLET_SYNC_SPAWN = s
    coach_api._DROPLET_PROVISION_SPAWN = p


@pytest.fixture(autouse=True)
def _reset_registry() -> Iterator[None]:
    """Ensure no in-flight runs leak between tests (concurrency tests need this)."""
    yield
    # Mark every busy run as terminal so the next test sees a clean slate.
    for run in coach_api._REGISTRY.list_runs():
        if run.status in ("pending", "running"):
            coach_api._REGISTRY.publish(
                run.id,
                _runs.StatusEvent(run_id=run.id, status="cancelled", message="test teardown"),
            )
            coach_api._REGISTRY.close_subscribers(run.id)


def _parse_sse(text: str) -> list[dict[str, Any]]:
    out: list[dict[str, Any]] = []
    for frame in text.split("\n\n"):
        if not frame.strip():
            continue
        kind = ""
        data = ""
        for ln in frame.splitlines():
            if ln.startswith("event: "):
                kind = ln[7:]
            elif ln.startswith("data: "):
                data = ln[6:]
        if kind and data:
            out.append({"event": kind, "data": json.loads(data)})
    return out


# ---- /api/github/status --------------------------------------------------


def test_github_status_unconfigured_lists_missing(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.delenv("GITHUB_TOKEN", raising=False)
    monkeypatch.delenv("GITHUB_REPO", raising=False)
    r = client.get("/coach/api/github/status")
    assert r.status_code == 200
    body = r.json()
    assert body["configured"] is False
    assert "GITHUB_TOKEN" in body["missing"]
    assert "GITHUB_REPO" in body["missing"]


def test_github_status_configured(monkeypatch: pytest.MonkeyPatch) -> None:
    for k, v in _GITHUB_ENV.items():
        monkeypatch.setenv(k, v)
    # Force "binaries present" by monkeypatching shutil.which for predictable CI.
    monkeypatch.setattr("shutil.which", lambda b: f"/usr/bin/{b}")
    r = client.get("/coach/api/github/status")
    assert r.status_code == 200
    body = r.json()
    assert body["configured"] is True
    assert body["missing"] == []
    assert body["target"] == "professor-codephreak/mindXtrain"


# ---- /api/github/push ----------------------------------------------------


def test_github_push_503_when_unconfigured(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.delenv("GITHUB_TOKEN", raising=False)
    monkeypatch.delenv("GITHUB_REPO", raising=False)
    r = client.post("/coach/api/github/push", json={})
    assert r.status_code == 503
    detail = r.json()["detail"]
    assert "missing" in detail
    assert "GITHUB_TOKEN" in detail["missing"]


def test_github_push_200_with_fake_spawn(monkeypatch: pytest.MonkeyPatch) -> None:
    for k, v in _GITHUB_ENV.items():
        monkeypatch.setenv(k, v)
    monkeypatch.setattr("shutil.which", lambda b: f"/usr/bin/{b}")

    captured: dict[str, str] = {}

    def _fake(run: _runs.Run, req: coach_api.GithubPushRequest) -> None:
        captured["run_id"] = run.id
        captured["msg"] = req.commit_message
        coach_api._REGISTRY.publish(
            run.id,
            _runs.LogEvent(run_id=run.id, line="=== step 1/11: probe-git ==="),
        )
        coach_api._REGISTRY.publish(
            run.id,
            _runs.StatusEvent(run_id=run.id, status="succeeded", message="rc=0"),
        )

    coach_api._GITHUB_SPAWN = _fake
    r = client.post("/coach/api/github/push", json={"commit_message": "hello"})
    assert r.status_code == 200
    body = r.json()
    assert body["recipe"] == "_github_push"
    assert body["id"] == captured["run_id"]
    assert captured["msg"] == "hello"


# ---- /api/droplet/status ------------------------------------------------


def test_droplet_status_returns_both_modes(monkeypatch: pytest.MonkeyPatch) -> None:
    for k, v in _DROPLET_ENV.items():
        monkeypatch.setenv(k, v)
    for k, v in _AMD_DC_ENV.items():
        monkeypatch.setenv(k, v)
    monkeypatch.setattr("shutil.which", lambda b: f"/usr/bin/{b}")
    r = client.get("/coach/api/droplet/status")
    assert r.status_code == 200
    body = r.json()
    assert "sync" in body and "provision" in body
    assert body["sync"]["configured"] is True
    assert body["provision"]["configured"] is True
    assert "mi300x.test" in body["sync"]["target"]
    assert body["provision"]["target"] == "amd-dev-cloud:atl1:gpu-mi300x8-1536gb-devcloud"


def test_droplet_status_unconfigured(monkeypatch: pytest.MonkeyPatch) -> None:
    for k in (*_DROPLET_ENV, *_AMD_DC_ENV):
        monkeypatch.delenv(k, raising=False)
    r = client.get("/coach/api/droplet/status")
    assert r.status_code == 200
    body = r.json()
    assert body["sync"]["configured"] is False
    assert body["provision"]["configured"] is False
    assert "DROPLET_HOST" in body["sync"]["missing"]
    assert "AMD_DEV_CLOUD_TOKEN" in body["provision"]["missing"]


# ---- /api/droplet/sync --------------------------------------------------


def test_droplet_sync_503_when_unconfigured(monkeypatch: pytest.MonkeyPatch) -> None:
    for k in _DROPLET_ENV:
        monkeypatch.delenv(k, raising=False)
    r = client.post("/coach/api/droplet/sync", json={})
    assert r.status_code == 503
    assert "DROPLET_HOST" in r.json()["detail"]["missing"]


def test_droplet_sync_409_when_provision_in_flight(monkeypatch: pytest.MonkeyPatch) -> None:
    for k, v in {**_DROPLET_ENV, **_AMD_DC_ENV}.items():
        monkeypatch.setenv(k, v)
    monkeypatch.setattr("shutil.which", lambda b: f"/usr/bin/{b}")

    # Spawn a provision that stays in-flight.
    def _hang_provision(run: _runs.Run, _req: coach_api.DropletProvisionRequest) -> None:
        coach_api._REGISTRY.publish(
            run.id, _runs.StatusEvent(run_id=run.id, status="running", message="pretending")
        )

    coach_api._DROPLET_PROVISION_SPAWN = _hang_provision
    r1 = client.post("/coach/api/droplet/provision", json={})
    assert r1.status_code == 200, r1.text

    # Now a sync should 409.
    r2 = client.post("/coach/api/droplet/sync", json={})
    assert r2.status_code == 409
    detail = r2.json()["detail"]
    assert detail["active_recipe"] == "_droplet_provision"


def test_droplet_sync_200_with_fake_spawn(monkeypatch: pytest.MonkeyPatch) -> None:
    for k, v in _DROPLET_ENV.items():
        monkeypatch.setenv(k, v)
    monkeypatch.setattr("shutil.which", lambda b: f"/usr/bin/{b}")

    def _fake(run: _runs.Run, req: coach_api.DropletSyncRequest) -> None:
        assert req.run_bench is True
        coach_api._REGISTRY.publish(
            run.id, _runs.LogEvent(run_id=run.id, line="=== step 1/4: rsync ===")
        )
        coach_api._REGISTRY.publish(
            run.id, _runs.StatusEvent(run_id=run.id, status="succeeded", message="rc=0")
        )

    coach_api._DROPLET_SYNC_SPAWN = _fake
    r = client.post("/coach/api/droplet/sync", json={})
    assert r.status_code == 200
    body = r.json()
    assert body["recipe"] == "_droplet_sync"


# ---- /api/droplet/provision ---------------------------------------------


def test_droplet_provision_503_when_unconfigured(monkeypatch: pytest.MonkeyPatch) -> None:
    for k in _AMD_DC_ENV:
        monkeypatch.delenv(k, raising=False)
    r = client.post("/coach/api/droplet/provision", json={})
    assert r.status_code == 503
    assert "AMD_DEV_CLOUD_TOKEN" in r.json()["detail"]["missing"]


def test_droplet_provision_passes_request_through_to_spawn(monkeypatch: pytest.MonkeyPatch) -> None:
    for k, v in _AMD_DC_ENV.items():
        monkeypatch.setenv(k, v)

    captured: dict[str, str] = {}

    def _fake(run: _runs.Run, req: coach_api.DropletProvisionRequest) -> None:
        captured["repo"] = req.repo or "default"
        captured["wait"] = str(req.wait_for_bootstrap)
        coach_api._REGISTRY.publish(
            run.id, _runs.StatusEvent(run_id=run.id, status="succeeded", message="rc=0")
        )

    coach_api._DROPLET_PROVISION_SPAWN = _fake
    r = client.post("/coach/api/droplet/provision", json={"repo": "owner/repo", "wait_for_bootstrap": False})
    assert r.status_code == 200
    assert captured["repo"] == "owner/repo"
    assert captured["wait"] == "False"


# ---- SSE replay over the deploy run -------------------------------------


def test_deploy_run_sse_replays_log_and_status(monkeypatch: pytest.MonkeyPatch) -> None:
    for k, v in _GITHUB_ENV.items():
        monkeypatch.setenv(k, v)
    monkeypatch.setattr("shutil.which", lambda b: f"/usr/bin/{b}")

    def _fake(run: _runs.Run, _req: coach_api.GithubPushRequest) -> None:
        coach_api._REGISTRY.publish(
            run.id, _runs.LogEvent(run_id=run.id, line="=== step 1/11 ===")
        )
        coach_api._REGISTRY.publish(
            run.id, _runs.LogEvent(run_id=run.id, line="initialized empty git repo")
        )
        coach_api._REGISTRY.publish(
            run.id, _runs.StatusEvent(run_id=run.id, status="succeeded", message="rc=0")
        )

    coach_api._GITHUB_SPAWN = _fake
    r = client.post("/coach/api/github/push", json={})
    run_id = r.json()["id"]

    # Give the registry a moment for the synchronous fake to publish.
    time.sleep(0.05)

    es = client.get(f"/coach/api/runs/{run_id}/events", headers={"accept": "text/event-stream"})
    assert es.status_code == 200
    events = _parse_sse(es.text)
    kinds = [e["event"] for e in events]
    assert "log" in kinds
    assert "status" in kinds
    final = [e for e in events if e["event"] == "status"][-1]
    assert final["data"]["status"] == "succeeded"


# ---- /api/runs/{id}/cancel works on synthetic runs ----------------------


def test_cancel_endpoint_works_on_deploy_run(monkeypatch: pytest.MonkeyPatch) -> None:
    for k, v in _GITHUB_ENV.items():
        monkeypatch.setenv(k, v)
    monkeypatch.setattr("shutil.which", lambda b: f"/usr/bin/{b}")

    def _fake(run: _runs.Run, _req: coach_api.GithubPushRequest) -> None:
        coach_api._REGISTRY.publish(
            run.id, _runs.StatusEvent(run_id=run.id, status="running", message="busy")
        )

    coach_api._GITHUB_SPAWN = _fake
    r = client.post("/coach/api/github/push", json={})
    run_id = r.json()["id"]

    c = client.post(f"/coach/api/runs/{run_id}/cancel")
    # No real subprocess means cancel returns False, but the endpoint still 200s.
    assert c.status_code == 200
    assert "cancelled" in c.json()