Coverage for tests/streamwise/test_streamwise.py: 100%
401 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 HTTP server routes.
4"""
6import sys
7import pytest
8from datetime import datetime, timezone
10from http import HTTPStatus
12from unittest.mock import patch
13from unittest.mock import AsyncMock
15from quart.typing import TestClientProtocol
17from streamwise import http_session_manager
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
30 from streamwise.service_account_manager import get_streamwise_service_account
31 from streamwise.service_account_manager import get_streamwiseapp_service_account
33 from streamwise.service_manager import get_k8s_container_logs
34 from streamwise.service_manager import get_k8s_pod_events
35 from streamwise.service_manager import get_services
36 from streamwise.service_manager import get_services_ns
37 from streamwise.service_manager import get_service_timestamps
38 from streamwise.service_manager import get_service_health
39 from streamwise.service_manager import get_service_files
40 from streamwise.service_manager import get_health_and_files_async
41 from streamwise.service_manager import parse_vllm_metrics
44def _get_client() -> TestClientProtocol:
45 app = sw.app
46 client = app.test_client()
47 return client
50@pytest.fixture(scope="function", autouse=True)
51def setup_k8s_cluster() -> None:
52 # for some reason k8s_config.load_kube_config() is not async mocked
53 sw.k8s_cluster = "unittest"
56@pytest.mark.asyncio
57async def test_index() -> None:
58 client = _get_client()
59 response = await client.get("/")
60 assert response.status_code == HTTPStatus.OK
61 response_text = await response.get_data(as_text=True)
62 assert "StreamWise Cluster Manager" in response_text
63 assert "Applications" in response_text
64 assert "Wrappers" in response_text
67@pytest.mark.asyncio
68async def test_index_incluster() -> None:
69 sw.k8s_cluster = "incluster"
71 client = _get_client()
72 response = await client.get("/")
73 assert response.status_code == HTTPStatus.OK
74 response_text = await response.get_data(as_text=True)
75 assert "StreamWise Cluster Manager" in response_text
76 assert "Applications" in response_text
77 assert "Wrappers" in response_text
80@pytest.mark.asyncio
81async def test_health() -> None:
82 sw.k8s_cluster = "clustername0"
84 client = _get_client()
85 response = await client.get("/health")
86 assert response.status_code == HTTPStatus.OK
87 response_json = await response.get_json()
88 assert response_json == {
89 "status": "ok",
90 "k8s_cluster": "clustername0"
91 }
94@pytest.mark.asyncio
95async def test_service_info() -> None:
96 sw.k8s_cluster = "unittest"
98 client = _get_client()
99 response = await client.get("/service/test")
100 assert response.status_code == HTTPStatus.OK
101 response_text = await response.get_data(as_text=True)
102 assert response_text.startswith("<!DOCTYPE html>\n<html lang=\"en\">")
103 assert "<title>test</title>" in response_text
106@pytest.mark.asyncio
107async def test_service_info_flux() -> None:
108 client = _get_client()
109 response = await client.get("/service/flux")
110 assert response.status_code == HTTPStatus.OK
111 response_text = await response.get_data(as_text=True)
112 assert response_text.startswith("<!DOCTYPE html>\n<html lang=\"en\">")
113 assert "<title>FLUX</title>" in response_text
116@pytest.mark.asyncio
117async def test_service_info_fantasytalking() -> None:
118 app = sw.app
119 client = app.test_client()
120 response = await client.get("/service/fantasytalking")
121 assert response.status_code == HTTPStatus.OK
122 response_text = await response.get_data(as_text=True)
123 assert response_text.startswith("<!DOCTYPE html>\n<html lang=\"en\">")
124 assert "<title>Fantasy Talking</title>" in response_text
127@pytest.mark.asyncio
128async def test_service_info_has_submit_job_button() -> None:
129 client = _get_client()
130 response = await client.get("/service/test")
131 assert response.status_code == HTTPStatus.OK
132 response_text = await response.get_data(as_text=True)
133 assert 'href="/job/"' in response_text
134 assert 'title="Submit Job"' in response_text
137_MOCK_MIG_SERVICE = {
138 "namespace": "rtgen",
139 "pod_name": "realesrgan-pod",
140 "pod_ip": "10.0.0.6",
141 "container_port": 8080,
142 "container_name": "realesrgan",
143 "pod_status": "Running",
144 "start_time": None,
145 "url": "http://10.0.0.6:8080",
146 "node_name": "testnode",
147 "cpu": 1,
148 "memory": 2 * 1024 * 1024 * 1024,
149 "gpu": 1,
150 "mig_profile": "1g.10gb",
151 "ephemeral_storage": 0,
152 "events": [],
153 "image": "myacr.azurecr.io/realesrgan:latest",
154 "logs": None,
155 "health": None,
156 "files": None,
157}
160@pytest.mark.asyncio
161async def test_service_shows_mig_profile() -> None:
162 """Service page shows MIG profile badge for MIG pods."""
163 client = _get_client()
164 with patch("streamwise.streamwise.get_services", new=AsyncMock(return_value=[_MOCK_MIG_SERVICE])):
165 with patch("streamwise.streamwise.get_k8s_load_balancers", new=AsyncMock(return_value=[])):
166 response = await client.get("/service/realesrgan")
167 assert response.status_code == HTTPStatus.OK
168 response_text = await response.get_data(as_text=True)
169 assert "1g.10gb" in response_text
170 assert "MIG slice" in response_text
173@pytest.mark.asyncio
174async def test_service_shows_full_gpu() -> None:
175 """Service page shows plain GPU count for full-GPU pods (no MIG)."""
176 client = _get_client()
177 full_gpu_svc = dict(_MOCK_MIG_SERVICE)
178 full_gpu_svc["mig_profile"] = None
179 full_gpu_svc["container_name"] = "flux"
180 with patch("streamwise.streamwise.get_services", new=AsyncMock(return_value=[full_gpu_svc])):
181 with patch("streamwise.streamwise.get_k8s_load_balancers", new=AsyncMock(return_value=[])):
182 response = await client.get("/service/flux")
183 assert response.status_code == HTTPStatus.OK
184 response_text = await response.get_data(as_text=True)
185 assert "MIG slice" not in response_text
188@pytest.mark.asyncio
189async def test_service_multiple_pods_no_duplicate_http_error_class() -> None:
190 """Service page with multiple pods must declare HttpError exactly once (no redeclaration error)."""
191 client = _get_client()
192 svc1 = dict(_MOCK_MIG_SERVICE)
193 svc1["mig_profile"] = None
194 svc1["pod_name"] = "realesrgan-pod-0"
195 svc1["pod_ip"] = "10.0.0.6"
196 svc2 = dict(_MOCK_MIG_SERVICE)
197 svc2["mig_profile"] = None
198 svc2["pod_name"] = "realesrgan-pod-1"
199 svc2["pod_ip"] = "10.0.0.7"
200 with patch("streamwise.streamwise.get_services", new=AsyncMock(return_value=[svc1, svc2])):
201 with patch("streamwise.streamwise.get_k8s_load_balancers", new=AsyncMock(return_value=[])):
202 response = await client.get("/service/realesrgan")
203 assert response.status_code == HTTPStatus.OK
204 response_text = await response.get_data(as_text=True)
205 assert response_text.count("class HttpError extends Error") == 1
208@pytest.mark.asyncio
209async def test_container_info() -> None:
210 app = sw.app
211 client = app.test_client()
212 response = await client.get("/service/test/10.1.1.1")
213 assert response.status_code == HTTPStatus.OK
214 response_text = await response.get_data(as_text=True)
215 assert response_text.startswith("<!DOCTYPE html>\n<html lang=\"en\">")
216 assert "<title>test</title>" in response_text
217 # assert "10.1.1.1" in response_text
220@pytest.mark.asyncio
221async def test_service_timeline() -> None:
222 client = _get_client()
223 response = await client.get("/service/test/timeline")
224 assert response.status_code == HTTPStatus.OK
225 response_text = await response.get_data(as_text=True)
226 assert response_text.startswith("<!DOCTYPE html>\n<html lang=\"en\">")
227 assert "<title>test timeline</title>" in response_text
230@pytest.mark.asyncio
231async def test_timeline() -> None:
232 client = _get_client()
233 response = await client.get("/service/timeline")
234 assert response.status_code == HTTPStatus.OK
235 response_json = await response.get_json()
236 assert response_json is None
239@pytest.mark.asyncio
240async def test_api_add_pod() -> None:
241 client = _get_client()
243 # POST without JSON should fail
244 response = await client.post("/api/pod", data={
245 "container_name": "flux",
246 })
247 assert response.status_code == HTTPStatus.BAD_REQUEST
248 response_json = await response.get_json()
249 assert response_json == {"error": "Missing required parameter 'container_name'"}
252@pytest.mark.asyncio
253async def test_nonexisting() -> None:
254 client = _get_client()
255 await client.get("/nonexisting")
258def test_parse_vllm_metrics() -> None:
259 sample_metrics = """
260vllm:num_requests_running{engine="0",model_name="google/gemma-3-27b-it"} 0.0
261vllm:request_success_total{engine="0",finished_reason="stop",model_name="google/gemma-3-27b-it"} 0.0
262vllm:request_prompt_tokens_bucket{engine="0",le="1.0",model_name="google/gemma-3-27b-it"} 0.0
263vllm:request_success_total{engine="0",finished_reason="stop",model_name="google/gemma-3-27b-it"} 0.0
264vllm:request_success_total{engine="0",finished_reason="length",model_name="google/gemma-3-27b-it"} 0.0
265vllm:request_success_total{engine="0",finished_reason="abort",model_name="google/gemma-3-27b-it"} 0.0
266"""
267 result = parse_vllm_metrics(sample_metrics)
268 assert len(result) == 3
269 assert result == {
270 "num_requests_running": 0.0,
271 "request_success_total": 0.0,
272 "request_prompt_tokens_bucket": 0.0,
273 }
276@pytest.mark.asyncio
277async def test_get_k8s_container_logs() -> None:
278 k8s_api_mock = AsyncMock()
279 k8s_api_mock.read_namespaced_pod_log = AsyncMock()
280 k8s_api_mock.read_namespaced_pod_log.return_value = "log line"
282 logs = await get_k8s_container_logs(
283 k8s_api=k8s_api_mock,
284 namespace="namespace",
285 pod_name="pod_name",
286 container_name="container_name")
287 assert logs == "log line"
290@pytest.mark.asyncio
291async def test_get_k8s_pod_events() -> None:
292 k8s_api_mock = AsyncMock()
294 events = await get_k8s_pod_events(
295 k8s_api=k8s_api_mock,
296 namespace="namespace",
297 pod_name="pod_name")
298 assert events == []
301@pytest.mark.asyncio
302async def test_get_services() -> None:
303 services = await get_services(k8s_cluster="unittest")
304 assert services == []
306 services = await get_services_ns(
307 namespace="namespace",
308 k8s_cluster="unittest")
309 assert services == []
312@pytest.mark.asyncio
313async def test_get_health_and_files_async() -> None:
314 service_data = await get_health_and_files_async(None) # type: ignore[arg-type]
315 assert service_data is None
317 service_data = await get_health_and_files_async([])
318 assert service_data == []
320 service_data = await get_health_and_files_async([
321 {"service1": []}
322 ])
323 assert service_data == [
324 {
325 "files": [],
326 "health": "N/A",
327 "service1": []
328 }
329 ]
332@pytest.mark.asyncio
333async def test_get_service_account() -> None:
334 sa = await get_streamwise_service_account(k8s_cluster="unittest")
335 assert sa == "streamwise-service-account"
337 sa = await get_streamwiseapp_service_account(k8s_cluster="unittest")
338 assert sa == "streamwiseapp-service-account"
341@pytest.mark.asyncio
342async def test_start_stop() -> None:
343 await sw.startup()
344 # TODO test http_session_manager.client_session
345 # assert http_session_manager.client_session is not None
346 # assert sw.client_session is not None
348 # session = await sw.get_global_session()
349 # assert session is not None
351 session = await http_session_manager.get_global_session()
352 assert session is not None
354 await sw.shutdown()
355 # assert sw.client_session is None
358@pytest.mark.asyncio
359async def test_get_service_timestamps() -> None:
360 timestamps = await get_service_timestamps(
361 pod_name="test-pod",
362 container_name="test-container",
363 url="http://10.1.2.3:8080")
364 assert timestamps is not None
365 assert timestamps == []
367 timestamps = await get_service_timestamps(
368 pod_name="test-pod",
369 container_name="test-container",
370 url=None) # type: ignore[arg-type]
371 assert timestamps is None
373 timestamps = await get_service_timestamps(
374 pod_name="test-pod",
375 container_name="gemma",
376 url="http://20.1.2.3:8888")
377 assert timestamps == []
380@pytest.mark.asyncio
381async def test_get_service_health() -> None:
382 health = await get_service_health(
383 container_name="test-container",
384 url="http://10.2.2.2:8080")
385 assert health is not None
386 assert health == {"status": "failed"}
388 health = await get_service_health(
389 container_name="test-container",
390 url=None) # type: ignore[arg-type]
391 assert health is None
394@pytest.mark.asyncio
395async def test_get_service_files() -> None:
396 files = await get_service_files(
397 container_name="test-container",
398 url="http://10.2.2.2:8080")
399 assert files == []
401 files = await get_service_files(
402 container_name="test-container",
403 url="N/A")
404 assert files is None
406 files = await get_service_files(
407 container_name="gemma",
408 url="http://10.2.2.2:8080")
409 assert files == []
412# ---------------------------------------------------------------------------
413# MIG-related UI fix tests
414# ---------------------------------------------------------------------------
416_MOCK_MIG_SERVICE_WITH_HEALTH = {
417 **_MOCK_MIG_SERVICE,
418 "health": {
419 "gpu": "NVIDIA A100-SXM4-80GB MIG 1g.10gb",
420 "world_size": 0
421 },
422}
424# Service with gpu_info where memory fields are None (e.g. MIG instance where
425# nvmlDeviceGetMemoryInfo raises NVMLError).
426_MOCK_MIG_SERVICE_WITH_NULL_MEM = {
427 **_MOCK_MIG_SERVICE,
428 "health": {
429 "gpu": "MIG 1g.10gb",
430 "world_size": 0,
431 "gpu_info": [
432 {
433 "index": 0,
434 "current": True,
435 "name": "MIG 1g.10gb",
436 "sm_util": None,
437 "mem_util": None,
438 "mem_gib_used": None,
439 "mem_gib_total": None,
440 "temp": None,
441 "power_draw_watts": None,
442 "power_limit_watts": None,
443 "graphics_clock": None,
444 "sm_clock": None,
445 "mem_clock": None,
446 }
447 ],
448 },
449}
451_MOCK_MIG_NODE = {
452 "node_name": "mig-node",
453 "region": "eastus",
454 "resource_group": "rg1",
455 "addresses": [],
456 "is_ready": True,
457 "capacity_resources": {
458 "cpu": 96,
459 "memory": 900 * 1024 * 1024 * 1024,
460 "storage": 500 * 1024 * 1024 * 1024,
461 "gpu": "N/A",
462 },
463 "allocatable_resources": {
464 "cpu": 96,
465 "memory": 900 * 1024 * 1024 * 1024,
466 "storage": 500 * 1024 * 1024 * 1024,
467 "gpu": "N/A",
468 },
469 "architecture": "amd64",
470 "kernel_version": "5.15.0",
471 "os_image": "Ubuntu 22.04",
472 "creation_timestamp": datetime(2024, 1, 1, tzinfo=timezone.utc),
473 "labels": {},
474 "images": [],
475 "gpu_model": "NVIDIA A100-SXM4-80GB",
476 "mig_enabled": True,
477 "mig_resources": {
478 "1g.10gb": {"capacity": 7, "allocatable": 5},
479 },
480}
482_MOCK_MIG_POD = {
483 "namespace": "rtgen",
484 "pod_name": "realesrgan-pod",
485 "status": "Running",
486 "pod_ip": "10.0.0.6",
487 "container_name": "realesrgan",
488 "url": "http://10.0.0.6:8080",
489 "node": "mig-node",
490 "cpu": 1,
491 "memory": 2 * 1024 * 1024 * 1024,
492 "gpu": 1,
493 "mig_profile": "1g.10gb",
494}
496_MOCK_FULL_GPU_NODE = {
497 **_MOCK_MIG_NODE,
498 "node_name": "gpu-node",
499 "gpu_model": "NVIDIA A100-SXM4-80GB",
500 "mig_enabled": False,
501 "mig_resources": {},
502 "capacity_resources": {
503 "cpu": 96,
504 "memory": 900 * 1024 * 1024 * 1024,
505 "storage": 500 * 1024 * 1024 * 1024,
506 "gpu": "8",
507 },
508 "allocatable_resources": {
509 "cpu": 96,
510 "memory": 900 * 1024 * 1024 * 1024,
511 "storage": 500 * 1024 * 1024 * 1024,
512 "gpu": "8",
513 },
514}
516_MOCK_FULL_GPU_POD = {
517 **_MOCK_MIG_POD,
518 "node": "gpu-node",
519 "mig_profile": None,
520 "gpu": 2,
521}
524@pytest.mark.asyncio
525async def test_service_shows_mig_with_gpu_model() -> None:
526 """Service page shows GPU model alongside MIG profile badge when health.gpu is available."""
527 client = _get_client()
528 with patch("streamwise.streamwise.get_services", new=AsyncMock(return_value=[_MOCK_MIG_SERVICE_WITH_HEALTH])):
529 with patch("streamwise.streamwise.get_k8s_load_balancers", new=AsyncMock(return_value=[])):
530 response = await client.get("/service/realesrgan")
531 assert response.status_code == HTTPStatus.OK
532 response_text = await response.get_data(as_text=True)
533 assert "1g.10gb" in response_text
534 assert "MIG slice" in response_text
535 # GPU model should appear alongside the MIG profile
536 assert "A100" in response_text
539@pytest.mark.asyncio
540async def test_service_mig_null_mem_renders_na() -> None:
541 """Service page renders N/A (not a TypeError) when mem_gib_used/total are None for MIG."""
542 client = _get_client()
543 with patch("streamwise.streamwise.get_services",
544 new=AsyncMock(return_value=[_MOCK_MIG_SERVICE_WITH_NULL_MEM])):
545 with patch("streamwise.streamwise.get_k8s_load_balancers", new=AsyncMock(return_value=[])):
546 response = await client.get("/service/realesrgan")
547 assert response.status_code == HTTPStatus.OK
548 response_text = await response.get_data(as_text=True)
549 assert "<td>N/A</td>" in response_text
550 assert "<th>Memory</th>" in response_text
553@pytest.mark.asyncio
554async def test_index_shows_mig_profile_in_wrappers() -> None:
555 """Index page wrappers table shows MIG profile badge instead of full GPU model for MIG services."""
556 client = _get_client()
557 with patch("streamwise.streamwise.get_services", new=AsyncMock(return_value=[_MOCK_MIG_SERVICE])):
558 with patch("streamwise.streamwise.get_k8s_nodes", new=AsyncMock(return_value=[])):
559 with patch("streamwise.streamwise.get_k8s_pods", new=AsyncMock(return_value=[])):
560 with patch("streamwise.streamwise.get_k8s_load_balancers", new=AsyncMock(return_value=[])):
561 response = await client.get("/")
562 assert response.status_code == HTTPStatus.OK
563 response_text = await response.get_data(as_text=True)
564 # MIG profile badge must appear in the wrappers table
565 assert "1g.10gb" in response_text
566 assert "MIG slice" in response_text
569@pytest.mark.asyncio
570async def test_index_shows_mig_profile_with_gpu_model_in_wrappers() -> None:
571 """Index page wrappers table shows GPU model alongside MIG profile when health.gpu is available."""
572 client = _get_client()
573 with patch("streamwise.streamwise.get_services", new=AsyncMock(return_value=[_MOCK_MIG_SERVICE_WITH_HEALTH])):
574 with patch("streamwise.streamwise.get_k8s_nodes", new=AsyncMock(return_value=[])):
575 with patch("streamwise.streamwise.get_k8s_pods", new=AsyncMock(return_value=[])):
576 with patch("streamwise.streamwise.get_k8s_load_balancers", new=AsyncMock(return_value=[])):
577 response = await client.get("/")
578 assert response.status_code == HTTPStatus.OK
579 response_text = await response.get_data(as_text=True)
580 assert "1g.10gb" in response_text
581 assert "MIG slice" in response_text
582 assert "A100" in response_text
585@pytest.mark.asyncio
586async def test_index_mig_node_shows_partitions() -> None:
587 """Index page nodes table shows MIG partitions (alloc/allocatable/capacity) for MIG-enabled nodes."""
588 client = _get_client()
589 with patch("streamwise.streamwise.get_services", new=AsyncMock(return_value=[])):
590 with patch("streamwise.streamwise.get_k8s_nodes", new=AsyncMock(return_value=[_MOCK_MIG_NODE])):
591 with patch("streamwise.streamwise.get_k8s_pods", new=AsyncMock(return_value=[_MOCK_MIG_POD])):
592 with patch("streamwise.streamwise.get_k8s_load_balancers", new=AsyncMock(return_value=[])):
593 response = await client.get("/")
594 assert response.status_code == HTTPStatus.OK
595 response_text = await response.get_data(as_text=True)
596 # MIG partition profile must appear in the nodes table
597 assert "<br>+1g.10gb: 1/5/7\n" in response_text
600@pytest.mark.asyncio
601async def test_index_mig_pods_excluded_from_full_gpu_count() -> None:
602 """Index page nodes table does not count MIG pods toward the full GPU allocated total."""
603 client = _get_client()
605 # Mixed node: has both full GPUs and MIG resources
606 mixed_node = {
607 **_MOCK_FULL_GPU_NODE,
608 "node_name": "mixed-node",
609 "mig_enabled": True,
610 "mig_resources": {
611 "1g.10gb": {
612 "capacity": 7,
613 "allocatable": 6
614 }
615 },
616 "capacity_resources": {
617 "cpu": 96,
618 "memory": 900 * 1024 * 1024 * 1024,
619 "storage": 500 * 1024 * 1024 * 1024,
620 "gpu": "7"
621 },
622 "allocatable_resources": {
623 "cpu": 96,
624 "memory": 900 * 1024 * 1024 * 1024,
625 "storage": 500 * 1024 * 1024 * 1024,
626 "gpu": "7"
627 },
628 }
629 # Two pods: one full-GPU pod (gpu=2, no mig_profile) and one MIG pod
630 full_pod = {**_MOCK_FULL_GPU_POD, "node": "mixed-node", "gpu": 2}
631 mig_pod = {**_MOCK_MIG_POD, "node": "mixed-node"}
632 with patch("streamwise.streamwise.get_services", new=AsyncMock(return_value=[])):
633 with patch("streamwise.streamwise.get_k8s_nodes", new=AsyncMock(return_value=[mixed_node])):
634 with patch("streamwise.streamwise.get_k8s_pods", new=AsyncMock(return_value=[full_pod, mig_pod])):
635 with patch("streamwise.streamwise.get_k8s_load_balancers", new=AsyncMock(return_value=[])):
636 response = await client.get("/")
637 assert response.status_code == HTTPStatus.OK
638 response_text = await response.get_data(as_text=True)
639 # The full GPU count should be 2 (only the non-MIG pod), not 3 (2+1 counting MIG pod)
640 assert "2/7/7" in response_text
641 assert "3/7/7" not in response_text
644def test_streamwise_arg_parser_https_args() -> None:
645 """streamwise.py argument parser accepts --certfile and --keyfile."""
646 import argparse
647 parser = argparse.ArgumentParser()
648 parser.add_argument("--k8s_cluster", type=str, default=None)
649 parser.add_argument("--host", type=str, default=sw.HOST)
650 parser.add_argument("--port", type=int, default=sw.PORT)
651 parser.add_argument("--certfile", type=str, default=None)
652 parser.add_argument("--keyfile", type=str, default=None)
653 parser.add_argument("--use-https", action="store_true", default=False)
654 args = parser.parse_args(["--certfile", "/tmp/cert.pem", "--keyfile", "/tmp/key.pem"])
655 assert args.certfile == "/tmp/cert.pem"
656 assert args.keyfile == "/tmp/key.pem"
657 assert args.use_https is False
660def test_streamwise_arg_parser_https_defaults() -> None:
661 """streamwise.py argument parser defaults certfile and keyfile to None."""
662 import argparse
663 parser = argparse.ArgumentParser()
664 parser.add_argument("--k8s_cluster", type=str, default=None)
665 parser.add_argument("--host", type=str, default=sw.HOST)
666 parser.add_argument("--port", type=int, default=sw.PORT)
667 parser.add_argument("--certfile", type=str, default=None)
668 parser.add_argument("--keyfile", type=str, default=None)
669 parser.add_argument("--use-https", action="store_true", default=False)
670 args = parser.parse_args([])
671 assert args.certfile is None
672 assert args.keyfile is None
673 assert args.use_https is False
676def test_set_service_scheme_https() -> None:
677 """set_service_scheme sets SERVICE_SCHEME to https."""
678 original = http_session_manager.SERVICE_SCHEME
679 try:
680 http_session_manager.set_service_scheme("https")
681 assert http_session_manager.SERVICE_SCHEME == "https"
682 finally:
683 http_session_manager.set_service_scheme(original)
686def test_set_service_scheme_http() -> None:
687 """set_service_scheme sets SERVICE_SCHEME back to http."""
688 original = http_session_manager.SERVICE_SCHEME
689 try:
690 http_session_manager.set_service_scheme("https")
691 http_session_manager.set_service_scheme("http")
692 assert http_session_manager.SERVICE_SCHEME == "http"
693 finally:
694 http_session_manager.set_service_scheme(original)
697def test_set_service_scheme_default() -> None:
698 """SERVICE_SCHEME defaults to http."""
699 assert http_session_manager.SERVICE_SCHEME == "http"
702def test_set_service_scheme_invalid() -> None:
703 """set_service_scheme raises ValueError for invalid schemes."""
704 import pytest as pt
705 with pt.raises(ValueError, match="Invalid service scheme"):
706 http_session_manager.set_service_scheme("ftp")
709def test_verify_ssl_default() -> None:
710 """VERIFY_SSL defaults to True."""
711 assert http_session_manager.VERIFY_SSL is True
714def test_set_verify_ssl_false() -> None:
715 """set_verify_ssl(False) disables SSL certificate verification."""
716 original = http_session_manager.VERIFY_SSL
717 try:
718 http_session_manager.set_verify_ssl(False)
719 assert http_session_manager.VERIFY_SSL is False
720 finally:
721 http_session_manager.set_verify_ssl(original)
724def test_set_verify_ssl_true() -> None:
725 """set_verify_ssl(True) re-enables SSL certificate verification."""
726 original = http_session_manager.VERIFY_SSL
727 try:
728 http_session_manager.set_verify_ssl(False)
729 http_session_manager.set_verify_ssl(True)
730 assert http_session_manager.VERIFY_SSL is True
731 finally:
732 http_session_manager.set_verify_ssl(original)
735@pytest.mark.asyncio
736async def test_api_cluster_gpus_aggregates_by_type() -> None:
737 """The cluster_gpus endpoint aggregates GPU counts by canonical type name."""
738 mock_nodes = [
739 {
740 "node_name": "h100-node-0",
741 "is_ready": True,
742 "gpu_model": "NVIDIA-H100-80GB-HBM3",
743 "allocatable_resources": {"gpu": "8"},
744 },
745 {
746 "node_name": "h100-node-1",
747 "is_ready": True,
748 "gpu_model": "NVIDIA-H100-80GB-HBM3",
749 "allocatable_resources": {"gpu": "8"},
750 },
751 {
752 "node_name": "a100-node-0",
753 "is_ready": True,
754 "gpu_model": "NVIDIA A100-SXM4-80GB",
755 "allocatable_resources": {"gpu": "8"},
756 },
757 {
758 "node_name": "cpu-node",
759 "is_ready": True,
760 "gpu_model": "N/A",
761 "allocatable_resources": {"gpu": "0"},
762 },
763 ]
764 client = _get_client()
765 with patch("streamwise.streamwise.get_k8s_nodes", new=AsyncMock(return_value=mock_nodes)):
766 response = await client.get("/api/auto_deploy/cluster_gpus")
767 assert response.status_code == HTTPStatus.OK
768 data = await response.get_json()
769 assert data["gpu_budget"] == {"H100": 16, "A100": 8}
772@pytest.mark.asyncio
773async def test_api_cluster_gpus_skips_not_ready_nodes() -> None:
774 """The cluster_gpus endpoint skips nodes that are not ready."""
775 mock_nodes = [
776 {
777 "node_name": "h100-node-0",
778 "is_ready": True,
779 "gpu_model": "NVIDIA-H100-80GB-HBM3",
780 "allocatable_resources": {"gpu": "8"},
781 },
782 {
783 "node_name": "h100-node-1",
784 "is_ready": False,
785 "gpu_model": "NVIDIA-H100-80GB-HBM3",
786 "allocatable_resources": {"gpu": "8"},
787 },
788 ]
789 client = _get_client()
790 with patch("streamwise.streamwise.get_k8s_nodes", new=AsyncMock(return_value=mock_nodes)):
791 response = await client.get("/api/auto_deploy/cluster_gpus")
792 assert response.status_code == HTTPStatus.OK
793 data = await response.get_json()
794 assert data["gpu_budget"] == {"H100": 8}
797@pytest.mark.asyncio
798async def test_api_cluster_gpus_empty_cluster() -> None:
799 """The cluster_gpus endpoint returns empty budget for a cluster with no GPU nodes."""
800 client = _get_client()
801 with patch("streamwise.streamwise.get_k8s_nodes", new=AsyncMock(return_value=[])):
802 response = await client.get("/api/auto_deploy/cluster_gpus")
803 assert response.status_code == HTTPStatus.OK
804 data = await response.get_json()
805 assert data["gpu_budget"] == {}