Coverage for tests/streamwise/test_streamwise_pod.py: 100%
240 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 pod_manager.py functions.
4"""
6import sys
7import pytest
8import urllib.parse
10from http import HTTPStatus
12from unittest.mock import patch, AsyncMock, MagicMock
14from tests.test_utils import temp_sys_path
15from tests.k8s_mock import K8sMock, MockApiException
17mock_k8s = K8sMock()
19mock_modules = {}
20mock_modules.update(mock_k8s.get_sub_modules())
21mock_k8s_client = mock_modules["kubernetes_asyncio.client"]
22mock_custom_api = mock_k8s_client.CustomObjectsApi.return_value
23with patch.dict(sys.modules, mock_modules):
24 with temp_sys_path("streamwise"):
25 from streamwise import streamwise as sw
27 from streamwise.pod_manager import get_gpu_type_affinity
28 from streamwise.pod_manager import get_container_port
29 from streamwise.pod_manager import get_gemma_settings
30 from streamwise.pod_manager import get_llama32_settings
31 from streamwise.pod_manager import get_mig_resource_name
32 from streamwise.pod_manager import get_tls_cert_settings
33 from streamwise.pod_manager import tls_cert_volume_exists
34 from streamwise.pod_manager import MIG_PROFILES
37@pytest.fixture(scope="function", autouse=True)
38def setup_k8s_cluster() -> None:
39 # for some reason k8s_config.load_kube_config() is not async mocked
40 sw.k8s_cluster = "unittest"
41 sw.use_https = False
44def test_get_gpu_type_affinity() -> None:
45 assert get_gpu_type_affinity("a100") == [
46 "NVIDIA-A100-SXM4-40GB",
47 "NVIDIA-A100-SXM4-80GB",
48 "NVIDIA-A100-PCIe-40GB",
49 "NVIDIA-A100-PCIe-80GB",
50 "NVIDIA-A100-80GB-PCIe",
51 ]
52 assert get_gpu_type_affinity("h100") == [
53 "NVIDIA-H100-SXM5-80GB",
54 "NVIDIA-H100-PCIe-80GB",
55 "NVIDIA-H100-NVL",
56 "NVIDIA-H100-80GB-HBM3",
57 "NVIDIA-H100",
58 ]
59 assert get_gpu_type_affinity("h200") == [
60 "NVIDIA-H200-SXM5-141GB",
61 "NVIDIA-H200",
62 ]
63 assert get_gpu_type_affinity("v100") == [
64 "Tesla-V100-PCIE-16GB",
65 "Tesla-V100-SXM2-16GB",
66 "Tesla-V100-SXM2-32GB"
67 ]
68 assert get_gpu_type_affinity("unknown") == []
69 assert get_gpu_type_affinity("a+") == [
70 "NVIDIA-A100-SXM4-40GB",
71 "NVIDIA-A100-SXM4-80GB",
72 "NVIDIA-A100-PCIe-40GB",
73 "NVIDIA-A100-PCIe-80GB",
74 "NVIDIA-A100-80GB-PCIe",
75 "NVIDIA-H100-SXM5-80GB",
76 "NVIDIA-H100-PCIe-80GB",
77 "NVIDIA-H100-NVL",
78 "NVIDIA-H100-80GB-HBM3",
79 "NVIDIA-H100",
80 "NVIDIA-H200-SXM5-141GB",
81 "NVIDIA-H200",
82 ]
85def test_get_container_port() -> None:
86 assert get_container_port("fantasytalking") == 8080
87 assert get_container_port("gemma") == 8000
88 assert get_container_port("streamwise") == 18181
89 assert get_container_port("streamcast") == 18080
90 assert get_container_port("flux") == 8080
93def test_get_gemma_settings() -> None:
94 args, volume_mounts, volumes = get_gemma_settings(num_gpus=1)
95 assert args is not None
96 assert "google/gemma-3-27b-it" in args
97 assert "1" in args
98 assert volume_mounts is not None
99 assert volumes is not None
102def test_get_llama32_settings() -> None:
103 args, volume_mounts, volumes = get_llama32_settings(num_gpus=2)
104 assert args is not None
105 assert "meta-llama/Llama-3.2-90B-Vision" in args
106 assert "--tensor-parallel-size" in args
107 assert "2" in args
108 assert volume_mounts is not None
109 assert volumes is not None
112def test_get_mig_resource_name() -> None:
113 # Valid A100 40 GB profiles
114 assert get_mig_resource_name("1g.5gb") == "nvidia.com/mig-1g.5gb"
115 assert get_mig_resource_name("2g.10gb") == "nvidia.com/mig-2g.10gb"
116 assert get_mig_resource_name("3g.20gb") == "nvidia.com/mig-3g.20gb"
117 assert get_mig_resource_name("4g.20gb") == "nvidia.com/mig-4g.20gb"
118 assert get_mig_resource_name("7g.40gb") == "nvidia.com/mig-7g.40gb"
119 # Valid A100 80 GB / H100 80 GB profiles
120 assert get_mig_resource_name("1g.10gb") == "nvidia.com/mig-1g.10gb"
121 assert get_mig_resource_name("2g.20gb") == "nvidia.com/mig-2g.20gb"
122 assert get_mig_resource_name("3g.40gb") == "nvidia.com/mig-3g.40gb"
123 assert get_mig_resource_name("4g.40gb") == "nvidia.com/mig-4g.40gb"
124 assert get_mig_resource_name("7g.80gb") == "nvidia.com/mig-7g.80gb"
125 # Invalid profile
126 assert get_mig_resource_name("invalid") is None
127 assert get_mig_resource_name("5g.20gb") is None
128 assert get_mig_resource_name("") is None
131def test_mig_profiles_set() -> None:
132 assert "1g.5gb" in MIG_PROFILES
133 assert "1g.10gb" in MIG_PROFILES
134 assert "7g.80gb" in MIG_PROFILES
135 assert "invalid" not in MIG_PROFILES
138def test_get_tls_cert_settings() -> None:
139 # Reset the relevant mocks so calls from earlier tests (e.g. get_gemma_settings)
140 # do not interfere with assert_called_once_with below.
141 mock_k8s_client.V1VolumeMount.reset_mock()
142 mock_k8s_client.V1Volume.reset_mock()
143 mock_k8s_client.V1CSIVolumeSource.reset_mock()
145 volume_mount, volume = get_tls_cert_settings()
146 assert volume_mount is not None
147 assert volume is not None
148 # V1VolumeMount / V1Volume are MagicMocks in the test environment; verify the
149 # constructors were called with the correct arguments.
150 mock_k8s_client.V1VolumeMount.assert_called_once_with(name="tls-csi", mount_path="/certs", read_only=True)
151 mock_k8s_client.V1CSIVolumeSource.assert_called_once_with(
152 driver="secrets-store.csi.k8s.io",
153 read_only=True,
154 volume_attributes={"secretProviderClass": "streamwise-tls"},
155 )
156 mock_k8s_client.V1Volume.assert_called_once_with(name="tls-csi", csi=mock_k8s_client.V1CSIVolumeSource.return_value)
159@pytest.mark.asyncio
160async def test_tls_cert_volume_exists_found() -> None:
161 """tls_cert_volume_exists returns True when the SecretProviderClass is found."""
162 mock_custom_api.get_namespaced_custom_object = AsyncMock(return_value={"metadata": {"name": "streamwise-tls"}})
163 result = await tls_cert_volume_exists("rtgen", "unittest")
164 assert result is True
165 mock_custom_api.get_namespaced_custom_object.assert_called_once_with(
166 group="secrets-store.csi.x-k8s.io",
167 version="v1",
168 namespace="rtgen",
169 plural="secretproviderclasses",
170 name="streamwise-tls"
171 )
174@pytest.mark.asyncio
175async def test_tls_cert_volume_exists_not_found() -> None:
176 """tls_cert_volume_exists returns False when the SecretProviderClass is absent (404)."""
177 mock_custom_api.get_namespaced_custom_object = AsyncMock(
178 side_effect=MockApiException(status=404, reason="Not Found")
179 )
180 result = await tls_cert_volume_exists("rtgen", "unittest")
181 assert result is False
184@pytest.mark.asyncio
185async def test_tls_cert_volume_exists_error() -> None:
186 """tls_cert_volume_exists returns False on unexpected errors."""
187 mock_custom_api.get_namespaced_custom_object = AsyncMock(
188 side_effect=Exception("connection refused")
189 )
190 result = await tls_cert_volume_exists("rtgen", "unittest")
191 assert result is False
194@pytest.mark.asyncio
195async def test_api_add_pod_no_tls_when_use_https_false() -> None:
196 """Pod creation without use_https=True must not call get_tls_cert_settings."""
197 sw.use_https = False
198 with patch.object(sw.pod_manager, "get_tls_cert_settings") as mock_get_tls, \
199 patch.object(sw.pod_manager, "tls_cert_volume_exists", new=AsyncMock(return_value=True)):
201 app = sw.app
202 client = app.test_client()
203 form_data = {"container_name": "qwenimageedit"}
204 response = await client.post(
205 "/api/pod",
206 data=urllib.parse.urlencode(form_data),
207 headers={"Content-Type": "application/x-www-form-urlencoded"},
208 )
209 assert response.status_code == HTTPStatus.OK
210 mock_get_tls.assert_not_called()
213@pytest.mark.asyncio
214async def test_api_add_pod_tls_added_when_use_https_and_volume_exists() -> None:
215 """Pod creation with use_https=True and existing SecretProviderClass must mount TLS volume."""
216 sw.use_https = True
217 with patch.object(sw.pod_manager, "get_tls_cert_settings") as mock_get_tls, \
218 patch.object(sw.pod_manager, "tls_cert_volume_exists", new=AsyncMock(return_value=True)):
219 mock_get_tls.return_value = (MagicMock(), MagicMock())
221 app = sw.app
222 client = app.test_client()
223 form_data = {"container_name": "qwenimageedit"}
224 response = await client.post(
225 "/api/pod",
226 data=urllib.parse.urlencode(form_data),
227 headers={"Content-Type": "application/x-www-form-urlencoded"},
228 )
229 assert response.status_code == HTTPStatus.OK
230 mock_get_tls.assert_called_once()
233@pytest.mark.asyncio
234async def test_api_add_pod_no_tls_when_volume_missing() -> None:
235 """Pod creation with use_https=True but missing SecretProviderClass must not mount TLS volume."""
236 sw.use_https = True
237 with patch.object(sw.pod_manager, "get_tls_cert_settings") as mock_get_tls, \
238 patch.object(sw.pod_manager, "tls_cert_volume_exists", new=AsyncMock(return_value=False)):
240 app = sw.app
241 client = app.test_client()
242 form_data = {"container_name": "qwenimageedit"}
243 response = await client.post(
244 "/api/pod",
245 data=urllib.parse.urlencode(form_data),
246 headers={"Content-Type": "application/x-www-form-urlencoded"},
247 )
248 assert response.status_code == HTTPStatus.OK
249 mock_get_tls.assert_not_called()
252@pytest.mark.asyncio
253async def test_add_pod() -> None:
254 app = sw.app
255 client = app.test_client()
256 response = await client.get("/pod/qwenimage")
257 assert response.status_code == HTTPStatus.OK
258 response_text = await response.get_data(as_text=True)
259 assert response_text.startswith("<!DOCTYPE html>\n<html lang=\"en\">")
260 assert "Add StreamWise Service" in response_text
263@pytest.mark.asyncio
264async def test_api_add_pod() -> None:
265 app = sw.app
266 client = app.test_client()
268 response = await client.post("/api/pod")
269 assert response.status_code == HTTPStatus.BAD_REQUEST
270 response_json = await response.get_json()
271 assert response_json == {"error": "Missing required parameter 'container_name'"}
273 # Actual content
274 form_data = {
275 "container_name": "qwenimageedit",
276 }
277 response = await client.post(
278 "/api/pod",
279 data=urllib.parse.urlencode(form_data),
280 headers={"Content-Type": "application/x-www-form-urlencoded"},
281 )
282 assert response.status_code == HTTPStatus.OK
283 response_json = await response.get_json()
284 assert response_json["container_name"] == "qwenimageedit"
285 assert response_json["message"] == "Pod creation requested"
286 assert response_json["pod_name"].startswith("qwenimageedit-")
287 assert "resource_request" in response_json
288 assert response_json["resource_request"] == {
289 "cpu": 2,
290 "ephemeral-storage": "16Gi",
291 "memory": "4Gi",
292 }
295@pytest.mark.asyncio
296async def test_api_add_pod_with_mig() -> None:
297 """Pod creation with a MIG profile should use the MIG resource name."""
298 app = sw.app
299 client = app.test_client()
301 form_data = {
302 "container_name": "kokoro",
303 "gpu": "1",
304 "mig_profile": "1g.5gb",
305 "gpu_type": "a100",
306 "memory": "8",
307 "cpu": "2",
308 }
309 response = await client.post(
310 "/api/pod",
311 data=urllib.parse.urlencode(form_data),
312 headers={"Content-Type": "application/x-www-form-urlencoded"},
313 )
314 assert response.status_code == HTTPStatus.OK
315 response_json = await response.get_json()
316 assert response_json["container_name"] == "kokoro"
317 assert response_json["mig_profile"] == "1g.5gb"
318 # Resource request must use MIG resource name, not nvidia.com/gpu
319 assert "nvidia.com/mig-1g.5gb" in response_json["resource_request"]
320 assert "nvidia.com/gpu" not in response_json["resource_request"]
323@pytest.mark.asyncio
324async def test_api_add_pod_invalid_mig() -> None:
325 """Pod creation with an invalid MIG profile should be rejected."""
326 app = sw.app
327 client = app.test_client()
329 form_data = {
330 "container_name": "kokoro",
331 "gpu": "1",
332 "mig_profile": "bad_profile",
333 }
334 response = await client.post(
335 "/api/pod",
336 data=urllib.parse.urlencode(form_data),
337 headers={"Content-Type": "application/x-www-form-urlencoded"},
338 )
339 assert response.status_code == HTTPStatus.BAD_REQUEST
340 response_json = await response.get_json()
341 assert "Invalid MIG profile" in response_json["error"]
344@pytest.mark.asyncio
345async def test_api_add_pod_custom_tag() -> None:
346 app = sw.app
347 client = app.test_client()
349 # Custom tag should be reflected in the image_url
350 form_data = {
351 "container_name": "flux",
352 "tag": "v9.9.9",
353 }
354 response = await client.post(
355 "/api/pod",
356 data=urllib.parse.urlencode(form_data),
357 headers={"Content-Type": "application/x-www-form-urlencoded"},
358 )
359 assert response.status_code == HTTPStatus.OK
360 response_json = await response.get_json()
361 assert response_json["container_name"] == "flux"
362 assert response_json["image_url"].endswith(":v9.9.9")
364 # Invalid container name should return 400 even with a tag
365 form_data_invalid = {
366 "container_name": "nonexistent-service",
367 "tag": "v1.0.0",
368 }
369 response = await client.post(
370 "/api/pod",
371 data=urllib.parse.urlencode(form_data_invalid),
372 headers={"Content-Type": "application/x-www-form-urlencoded"},
373 )
374 assert response.status_code == HTTPStatus.BAD_REQUEST
375 response_json = await response.get_json()
376 assert "error" in response_json
377 assert "nonexistent-service" in response_json["error"]
380@pytest.mark.asyncio
381async def test_remove_pod() -> None:
382 app = sw.app
383 client = app.test_client()
385 response = await client.delete("/api/pod/fluxkrea")
386 assert response.status_code == HTTPStatus.BAD_REQUEST
387 response_json = await response.get_json()
388 assert response_json == {"error": "Namespace is required"}
390 response = await client.delete("/api/pod/fluxkrea?namespace=default")
391 assert response.status_code == HTTPStatus.OK
392 response_json = await response.get_json()
393 assert response_json == {"message": "Pod fluxkrea removed successfully"}
396@pytest.mark.asyncio
397async def test_api_add_apps_success() -> None:
398 """POST /api/apps deploys all 9 application pods and returns 200."""
399 with patch.object(sw.pod_manager, "add_pod", new=AsyncMock(return_value=(
400 {"message": "Pod creation requested"}, HTTPStatus.OK
401 ))) as mock_add_pod:
402 app = sw.app
403 client = app.test_client()
405 response = await client.post("/api/apps")
406 assert response.status_code == HTTPStatus.OK
407 response_json = await response.get_json()
408 assert response_json == {"message": "Applications added successfully"}
410 # Verify add_pod was called once per application (9 apps total)
411 assert mock_add_pod.call_count == 9
413 # Verify each application received a unique lb_port
414 lb_ports = [call.kwargs["lb_port"] for call in mock_add_pod.call_args_list]
415 assert len(set(lb_ports)) == 9, "Each app must receive a unique lb_port"
417 # Verify all apps were deployed
418 app_names = [call.args[0] for call in mock_add_pod.call_args_list]
419 for expected in [
420 "streamcast", "streampersona", "streamchat", "streamshort",
421 "streammovie", "streamanimate", "streamlecture", "streamdub", "streamedit",
422 ]:
423 assert expected in app_names, f"{expected} was not deployed"
426@pytest.mark.asyncio
427async def test_api_add_apps_k8s_api_exception() -> None:
428 """POST /api/apps returns 500 when K8s raises an ApiException."""
429 with patch.object(sw.pod_manager, "add_pod", new=AsyncMock(
430 side_effect=MockApiException(
431 status=500,
432 reason="Internal Server Error",
433 body='{"message": "namespaces \\"rtgen\\" not found"}',
434 )
435 )):
436 app = sw.app
437 client = app.test_client()
439 response = await client.post("/api/apps")
440 assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
441 response_json = await response.get_json()
442 assert "error" in response_json
443 assert "rtgen" in response_json["error"]
446@pytest.mark.asyncio
447async def test_api_add_apps_generic_exception() -> None:
448 """POST /api/apps returns 500 when an unexpected exception is raised."""
449 with patch.object(sw.pod_manager, "add_pod", new=AsyncMock(
450 side_effect=Exception("unexpected failure")
451 )):
452 app = sw.app
453 client = app.test_client()
455 response = await client.post("/api/apps")
456 assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR
457 response_json = await response.get_json()
458 assert response_json == {"error": "unexpected failure"}