Coverage for tests/streamwise/test_streamwise_node.py: 100%
172 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1#!/usr/bin/env python3
2"""
3Unit tests for streamwise.py node-related endpoints.
4"""
6import sys
7import pytest
9from typing import Any
10from typing import Dict
12from unittest.mock import patch
13from unittest.mock import AsyncMock
15from quart.typing import TestClientProtocol
17from http import HTTPStatus
19from tests.test_utils import temp_sys_path
20from tests.k8s_mock import K8sMock
22mock_k8s = K8sMock()
24mock_modules = {}
25mock_modules.update(mock_k8s.get_sub_modules())
26with patch.dict(sys.modules, mock_modules):
27 with temp_sys_path("streamwise"):
28 from streamwise import streamwise as sw
29 # node_manager is imported by streamwise.py at module level;
30 # access it via the module attribute so we get the exact same object
31 node_manager = sw.node_manager
34def _get_client() -> TestClientProtocol:
35 app = sw.app
36 client = app.test_client()
37 return client
40@pytest.fixture(scope="function", autouse=True)
41def setup_k8s_cluster() -> None:
42 # for some reason k8s_config.load_kube_config() is not async mocked
43 sw.k8s_cluster = "unittest"
46@pytest.mark.asyncio
47async def test_nodes() -> None:
48 client = _get_client()
49 response = await client.get("/nodes")
50 assert response.status_code == HTTPStatus.NOT_FOUND
51 response_json = await response.get_json()
52 assert response_json == {"error": "No nodes found"}
55@pytest.mark.asyncio
56async def test_api_nodes() -> None:
57 client = _get_client()
58 response = await client.get("/api/nodes")
59 assert response.status_code == HTTPStatus.OK
60 response_json = await response.get_json()
61 assert response_json == []
64@pytest.mark.asyncio
65async def test_remove_node() -> None:
66 client = _get_client()
67 response = await client.delete("/api/node/testnode")
68 assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
69 response_json = await response.get_json()
70 assert response_json == {"error": "object MagicMock can't be used in 'await' expression"} # TODO
73@pytest.mark.asyncio
74async def test_node_info() -> None:
75 client = _get_client()
76 response = await client.get("/node/testnode")
77 assert response.status_code == HTTPStatus.NOT_FOUND
78 response_json = await response.get_json()
79 assert response_json == {"error": "Node 'testnode' not found"}
82_MOCK_NODE = {
83 "node_name": "testnode",
84 "region": "eastus",
85 "resource_group": "rg-test",
86 "addresses": [{
87 "type": "InternalIP",
88 "address": "10.0.0.1"
89 }],
90 "is_ready": True,
91 "capacity_resources": {
92 "cpu": 4.0,
93 "memory": 8 * 1024 * 1024 * 1024,
94 "storage": 100 * 1024 * 1024 * 1024,
95 "gpu": "N/A"
96 },
97 "allocatable_resources": {
98 "cpu": 4.0,
99 "memory": 8 * 1024 * 1024 * 1024,
100 "storage": 100 * 1024 * 1024 * 1024,
101 "gpu": "N/A"
102 },
103 "architecture": "amd64",
104 "kernel_version": "5.15.0",
105 "os_image": "Ubuntu 22.04",
106 "creation_timestamp": "2024-01-01T00:00:00Z",
107 "labels": {},
108 "images": None,
109 "gpu_model": "N/A",
110 "mig_enabled": False,
111 "mig_resources": {},
112}
115@pytest.mark.asyncio
116async def test_nodes_has_submit_job_button() -> None:
117 client = _get_client()
118 with patch.object(node_manager, "get_k8s_nodes", new=AsyncMock(return_value=[_MOCK_NODE])):
119 with patch.object(node_manager, "get_k8s_pods", new=AsyncMock(return_value=[])):
120 response = await client.get("/nodes")
121 assert response.status_code == HTTPStatus.OK
122 response_text = await response.get_data(as_text=True)
123 assert 'href="/job/"' in response_text
124 assert 'title="Submit Job"' in response_text
127@pytest.mark.asyncio
128async def test_node_shows_mig_enabled() -> None:
129 """Node page shows MIG enabled badge when mig_enabled is True."""
130 client = _get_client()
132 mig_node: Dict[str, Any] = dict(_MOCK_NODE)
133 mig_node.update({
134 "mig_enabled": True,
135 "gpu_model": "NVIDIA-A100-SXM4-80GB",
136 "capacity_resources": dict(mig_node["capacity_resources"], gpu=7),
137 "allocatable_resources": dict(mig_node["allocatable_resources"], gpu="N/A"),
138 "mig_resources": {
139 "1g.10gb": {
140 "capacity": 7,
141 "allocatable": 7
142 }
143 }
144 })
146 with patch.object(node_manager, "get_k8s_nodes", new=AsyncMock(return_value=[mig_node])):
147 with patch.object(node_manager, "get_k8s_pods", new=AsyncMock(return_value=[])):
148 response = await client.get("/nodes")
150 assert response.status_code == HTTPStatus.OK
151 response_text = await response.get_data(as_text=True)
152 msg = '<span class="text-success" title="Multi-Instance GPU" aria-label="MIG enabled">\u2705</span>'
153 assert msg in response_text
156@pytest.mark.asyncio
157async def test_node_shows_mig_resources() -> None:
158 """Node page shows per-profile MIG resource counts in the Resources table."""
159 client = _get_client()
161 mig_node: Dict[str, Any] = dict(_MOCK_NODE)
162 mig_node.update({
163 "mig_enabled": True,
164 "gpu_model": "NVIDIA-A100-SXM4-80GB",
165 "mig_resources": {
166 "1g.5gb": {
167 "capacity": 7,
168 "allocatable": 5
169 },
170 "2g.10gb": {
171 "capacity": 3,
172 "allocatable": 2
173 },
174 }
175 })
177 with patch.object(node_manager, "get_k8s_nodes", new=AsyncMock(return_value=[mig_node])):
178 with patch.object(node_manager, "get_k8s_pods", new=AsyncMock(return_value=[])):
179 response = await client.get("/nodes")
181 assert response.status_code == HTTPStatus.OK
182 response_text = await response.get_data(as_text=True)
183 assert "1g.5gb" in response_text
184 assert "2g.10gb" in response_text
185 # Allocatable and capacity counts visible
186 assert "5" in response_text
187 assert "2" in response_text
190@pytest.mark.asyncio
191async def test_node_shows_mig_disabled() -> None:
192 """Node page shows no MIG indicator when mig_enabled is False."""
193 client = _get_client()
194 with patch.object(node_manager, "get_k8s_nodes", new=AsyncMock(return_value=[_MOCK_NODE])):
195 with patch.object(node_manager, "get_k8s_pods", new=AsyncMock(return_value=[])):
196 response = await client.get("/nodes")
198 assert response.status_code == HTTPStatus.OK
199 response_text = await response.get_data(as_text=True)
200 assert "MIG" in response_text # row label is always present
201 assert "Enabled" not in response_text # but not the "Enabled" status text
204@pytest.mark.asyncio
205async def test_pod_shows_mig_profile() -> None:
206 """Node page shows MIG profile badge in the pods table for MIG pods."""
207 client = _get_client()
208 mig_pod = {
209 "namespace": "rtgen",
210 "pod_name": "flux-pod",
211 "status": "Running",
212 "pod_ip": "10.0.0.5",
213 "container_name": "flux",
214 "url": "http://10.0.0.5:8080",
215 "node": "testnode",
216 "cpu": 2,
217 "memory": 4 * 1024 * 1024 * 1024,
218 "gpu": 1,
219 "mig_profile": "1g.10gb",
220 }
221 with patch.object(node_manager, "get_k8s_nodes", new=AsyncMock(return_value=[_MOCK_NODE])):
222 with patch.object(node_manager, "get_k8s_pods", new=AsyncMock(return_value=[mig_pod])):
223 response = await client.get("/nodes")
225 assert response.status_code == HTTPStatus.OK
226 response_text = await response.get_data(as_text=True)
227 assert "1g.10gb" in response_text
228 assert "MIG slice" in response_text
231@pytest.mark.asyncio
232async def test_pod_shows_no_mig_profile_for_full_gpu() -> None:
233 """Node page shows plain GPU count for full-GPU pods (no MIG profile)."""
234 client = _get_client()
235 full_gpu_pod = {
236 "namespace": "rtgen",
237 "pod_name": "flux-pod",
238 "status": "Running",
239 "pod_ip": "10.0.0.5",
240 "container_name": "flux",
241 "url": "http://10.0.0.5:8080",
242 "node": "testnode",
243 "cpu": 2,
244 "memory": 4 * 1024 * 1024 * 1024,
245 "gpu": 1,
246 "mig_profile": None,
247 }
249 with patch.object(node_manager, "get_k8s_nodes", new=AsyncMock(return_value=[_MOCK_NODE])):
250 with patch.object(node_manager, "get_k8s_pods", new=AsyncMock(return_value=[full_gpu_pod])):
251 response = await client.get("/nodes")
253 assert response.status_code == HTTPStatus.OK
254 response_text = await response.get_data(as_text=True)
255 assert "MIG slice" not in response_text
258@pytest.mark.asyncio
259async def test_mig_visible_when_gpu_capacity_zero() -> None:
260 """Nodes table shows MIG profiles even when capacity_resources.gpu is 0 (single strategy)."""
261 client = _get_client()
262 mig_node: Dict[str, Any] = dict(_MOCK_NODE)
263 mig_node.update({
264 "mig_enabled": True,
265 "gpu_model": "NVIDIA-A100-SXM4-80GB",
266 "capacity_resources": dict(mig_node["capacity_resources"], gpu=0),
267 "allocatable_resources": dict(mig_node["allocatable_resources"], gpu=0),
268 "mig_resources": {
269 "1g.10gb": {"capacity": 3, "allocatable": 3},
270 "2g.20gb": {"capacity": 2, "allocatable": 2},
271 }
272 })
274 with patch.object(node_manager, "get_k8s_nodes", new=AsyncMock(return_value=[mig_node])):
275 with patch.object(node_manager, "get_k8s_pods", new=AsyncMock(return_value=[])):
276 response = await client.get("/nodes")
278 assert response.status_code == HTTPStatus.OK
279 response_text = await response.get_data(as_text=True)
280 assert "MIG" in response_text
281 assert "1g.10gb" in response_text
282 assert "2g.20gb" in response_text
285@pytest.mark.asyncio
286async def test_gpu_row_excludes_mig_pods() -> None:
287 """Resources table GPU row does not count MIG pods — they go in their own MIG rows."""
288 client = _get_client()
290 mig_node: Dict[str, Any] = dict(_MOCK_NODE)
291 mig_node.update({
292 "mig_enabled": True,
293 "gpu_model": "NVIDIA-A100-SXM4-80GB",
294 "capacity_resources": dict(mig_node["capacity_resources"], gpu=1),
295 "allocatable_resources": dict(mig_node["allocatable_resources"], gpu=1),
296 "mig_resources": {
297 "1g.10gb": {"capacity": 7, "allocatable": 7}
298 }
299 })
301 mig_pod = {
302 "namespace": "rtgen",
303 "pod_name": "esrgan-pod",
304 "status": "Running",
305 "pod_ip": "10.0.0.6",
306 "container_name": "realesrgan",
307 "url": "http://10.0.0.6:8080",
308 "node": "testnode",
309 "cpu": 1,
310 "memory": 2 * 1024 * 1024 * 1024,
311 "gpu": 1,
312 "mig_profile": "1g.10gb",
313 }
314 full_gpu_pod = {
315 "namespace": "rtgen",
316 "pod_name": "flux-pod",
317 "status": "Running",
318 "pod_ip": "10.0.0.5",
319 "container_name": "flux",
320 "url": "http://10.0.0.5:8080",
321 "node": "testnode",
322 "cpu": 2,
323 "memory": 4 * 1024 * 1024 * 1024,
324 "gpu": 1,
325 "mig_profile": None,
326 }
327 with patch.object(node_manager, "get_k8s_nodes", new=AsyncMock(return_value=[mig_node])):
328 with patch.object(node_manager, "get_k8s_pods", new=AsyncMock(return_value=[mig_pod, full_gpu_pod])):
329 response = await client.get("/node/testnode")
330 assert response.status_code == HTTPStatus.OK
331 response_text = await response.get_data(as_text=True)
332 # The GPU resource row should show "1.0" allocated (only the full-GPU pod),
333 # not "2.0" (which would include the MIG pod).
334 # Check the MIG row shows 1 allocated for 1g.10gb
335 assert "1g.10gb" in response_text
338@pytest.mark.asyncio
339async def test_mixed_full_gpu_and_mig() -> None:
340 """Nodes table shows both the full GPU count and MIG profiles (7 full + 1 MIG-partitioned)."""
341 client = _get_client()
342 mig_node: Dict[str, Any] = dict(_MOCK_NODE)
343 mig_node.update({
344 "mig_enabled": True,
345 "gpu_model": "NVIDIA-A100-SXM4-80GB",
346 "capacity_resources": dict(mig_node["capacity_resources"], gpu=7),
347 "allocatable_resources": dict(mig_node["allocatable_resources"], gpu=7),
348 "mig_resources": {
349 "1g.10gb": {"capacity": 3, "allocatable": 3},
350 "2g.20gb": {"capacity": 2, "allocatable": 2},
351 }
352 })
354 with patch.object(node_manager, "get_k8s_nodes", new=AsyncMock(return_value=[mig_node])):
355 with patch.object(node_manager, "get_k8s_pods", new=AsyncMock(return_value=[])):
356 response = await client.get("/nodes")
358 assert response.status_code == HTTPStatus.OK
359 response_text = await response.get_data(as_text=True)
360 assert "7" in response_text # Full GPU count visible
361 # MIG and profiles visible
362 assert "MIG" in response_text
363 assert "1g.10gb" in response_text
364 assert "2g.20gb" in response_text
367@pytest.mark.asyncio
368async def test_pure_mig_node_gpu_na() -> None:
369 """Resources table handles gpu=N/A without error on pure-MIG nodes."""
370 client = _get_client()
372 mig_node: Dict[str, Any] = dict(_MOCK_NODE)
373 mig_node.update({
374 "mig_enabled": True,
375 "gpu_model": "NVIDIA-A100-SXM4-80GB",
376 "capacity_resources": dict(mig_node["capacity_resources"], gpu="N/A"),
377 "allocatable_resources": dict(mig_node["allocatable_resources"], gpu="N/A"),
378 "mig_resources": {
379 "1g.10gb": {"capacity": 7, "allocatable": 7}
380 }
381 })
383 with patch.object(node_manager, "get_k8s_nodes", new=AsyncMock(return_value=[mig_node])):
384 with patch.object(node_manager, "get_k8s_pods", new=AsyncMock(return_value=[])):
385 response = await client.get("/node/testnode")
387 assert response.status_code == HTTPStatus.OK
388 response_text = await response.get_data(as_text=True)
389 assert "1g.10gb" in response_text
390 assert "MIG" in response_text