File size: 6,691 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 | """AMD Dev Cloud REST client — stubs the httpx transport, no network.
`httpx.MockTransport` is built into httpx (no `respx` dep needed).
"""
from __future__ import annotations
import json
from typing import Any
import httpx
import pytest
from mindxtrain.deploy.amd_dev_cloud import (
DEFAULT_BASE,
AmdDevCloudAuthError,
AmdDevCloudClient,
AmdDevCloudConfig,
AmdDevCloudError,
build_create_payload,
extract_public_ip,
from_env,
missing_env,
)
def _cfg(**ov: Any) -> AmdDevCloudConfig:
base = {"token": "dop_v1_TEST", "ssh_key_id": 56216059}
base.update(ov)
return AmdDevCloudConfig(**base) # type: ignore[arg-type]
def _client_with_handler(handler, *, cfg: AmdDevCloudConfig | None = None) -> AmdDevCloudClient:
cfg = cfg or _cfg()
transport = httpx.MockTransport(handler)
return AmdDevCloudClient(
cfg,
client=httpx.Client(
base_url=cfg.api_base,
headers={"Authorization": f"Bearer {cfg.token}"},
transport=transport,
),
)
def test_missing_env_lists_required() -> None:
miss = missing_env({})
assert "AMD_DEV_CLOUD_TOKEN" in miss
assert "AMD_DEV_CLOUD_SSH_KEY_ID" in miss
full = missing_env({"AMD_DEV_CLOUD_TOKEN": "x", "AMD_DEV_CLOUD_SSH_KEY_ID": "1"})
assert full == []
def test_from_env_parses_int_ssh_key_id() -> None:
cfg = from_env({
"AMD_DEV_CLOUD_TOKEN": "x",
"AMD_DEV_CLOUD_SSH_KEY_ID": "56216059",
"AMD_DEV_CLOUD_TAGS": "a, b ,c",
})
assert cfg.ssh_key_id == 56216059
assert cfg.tags == ("a", "b", "c")
assert cfg.api_base == DEFAULT_BASE
def test_from_env_rejects_non_integer_ssh_key_id() -> None:
with pytest.raises(RuntimeError, match="must be an int"):
from_env({"AMD_DEV_CLOUD_TOKEN": "x", "AMD_DEV_CLOUD_SSH_KEY_ID": "notanint"})
def test_build_create_payload_matches_curl_example() -> None:
"""Pin the exact JSON body the user supplied in the spec."""
cfg = _cfg(
region="atl1",
size="gpu-mi300x8-1536gb-devcloud",
image="vllm-0-17-1",
ssh_key_id=56216059,
tags=("mindx", "train", "aglm", "agenticplace", "pythai"),
)
body = build_create_payload(cfg, name="mindxtrain", user_data="")
# Field-by-field match to the curl example to lock the contract.
assert body == {
"name": "mindxtrain",
"region": "atl1",
"size": "gpu-mi300x8-1536gb-devcloud",
"image": "vllm-0-17-1",
"ssh_keys": [56216059],
"backups": False,
"ipv6": True,
"monitoring": True,
"tags": ["mindx", "train", "aglm", "agenticplace", "pythai"],
"user_data": "",
"vpc_uuid": "",
}
def test_create_returns_droplet_dict() -> None:
captured = {}
def handler(request: httpx.Request) -> httpx.Response:
assert request.method == "POST"
assert request.url.path == "/v2/droplets"
captured["auth"] = request.headers.get("authorization")
captured["body"] = json.loads(request.content)
return httpx.Response(202, json={"droplet": {"id": 12345, "status": "new", "name": "mindxtrain"}})
with _client_with_handler(handler) as c:
droplet = c.create(name="mindxtrain", user_data="#cloud-config\n")
assert droplet["id"] == 12345
assert captured["auth"] == "Bearer dop_v1_TEST"
assert captured["body"]["name"] == "mindxtrain"
assert captured["body"]["user_data"] == "#cloud-config\n"
def test_create_401_raises_auth_error() -> None:
def handler(_request: httpx.Request) -> httpx.Response:
return httpx.Response(401, json={"message": "bad token"})
with _client_with_handler(handler) as c:
with pytest.raises(AmdDevCloudAuthError, match="auth failed"):
c.create()
def test_create_500_raises_generic_error() -> None:
def handler(_request: httpx.Request) -> httpx.Response:
return httpx.Response(500, json={"message": "boom"})
with _client_with_handler(handler) as c:
with pytest.raises(AmdDevCloudError, match="500"):
c.create()
def test_poll_until_active_returns_when_status_active() -> None:
statuses = iter(["new", "new", "active"])
def handler(_request: httpx.Request) -> httpx.Response:
s = next(statuses)
body = {"droplet": {
"id": 1, "status": s,
"networks": {"v4": [{"type": "public", "ip_address": "1.2.3.4"}]},
}}
return httpx.Response(200, json=body)
sleeps: list[float] = []
with _client_with_handler(handler) as c:
droplet = c.poll_until_active(1, sleep=sleeps.append, interval=5.0, timeout=600)
assert droplet["status"] == "active"
# Three GETs → two sleeps between them.
assert sleeps == [5.0, 5.0]
def test_poll_until_active_raises_on_timeout() -> None:
def handler(_request: httpx.Request) -> httpx.Response:
return httpx.Response(200, json={"droplet": {"id": 1, "status": "new"}})
# Fake a clock that's always past the deadline after one tick.
times = iter([0.0, 0.0, 1000.0, 1000.0, 1000.0])
with _client_with_handler(handler) as c:
with pytest.raises(TimeoutError, match="did not reach 'active'"):
c.poll_until_active(1, timeout=10.0, interval=1.0, sleep=lambda _s: None, now=lambda: next(times))
def test_poll_until_active_raises_on_terminal_state() -> None:
def handler(_request: httpx.Request) -> httpx.Response:
return httpx.Response(200, json={"droplet": {"id": 1, "status": "errored"}})
with _client_with_handler(handler) as c:
with pytest.raises(AmdDevCloudError, match="terminal state"):
c.poll_until_active(1, sleep=lambda _s: None, now=lambda: 0.0, timeout=60)
def test_extract_public_ip_handles_no_v4() -> None:
assert extract_public_ip({"networks": {"v4": []}}) is None
assert extract_public_ip({"networks": {}}) is None
assert extract_public_ip({}) is None
def test_extract_public_ip_prefers_public_type() -> None:
droplet = {"networks": {"v4": [
{"type": "private", "ip_address": "10.0.0.1"},
{"type": "public", "ip_address": "1.2.3.4"},
]}}
assert extract_public_ip(droplet) == "1.2.3.4"
def test_list_filters_by_name() -> None:
def handler(_request: httpx.Request) -> httpx.Response:
return httpx.Response(200, json={"droplets": [
{"id": 1, "name": "mindxtrain"},
{"id": 2, "name": "other"},
]})
with _client_with_handler(handler) as c:
all_ = c.list()
only = c.list(name="mindxtrain")
assert len(all_) == 2
assert len(only) == 1
assert only[0]["id"] == 1
|