Coverage for tests/streamwise/test_container_config.py: 100%

1"""Tests for shared StreamWise container configuration.""" 

2 

3from __future__ import annotations 

4 

5from tests.test_utils import temp_sys_path 

6 

7with temp_sys_path("streamwise", "simulator"): 

8 from container_config import CONTAINER_RESOURCES 

9 from container_config import MODEL_CONTAINER_RESOURCES 

10 from container_config import MODEL_TO_CONTAINER_NAME 

11 from container_config import get_minimum_service_container_specs 

12 

13 

14def test_minimum_service_specs_include_expected_gpu_defaults() -> None: 

15 """GPU defaults should preserve current minimal /api/service behavior.""" 

16 specs = get_minimum_service_container_specs(max_gpus=8) 

17 

18 for container_name in CONTAINER_RESOURCES: 

19 assert container_name in specs 

20 

21 assert specs["gemma"].gpu == 2 

22 assert specs["kokoro"].gpu == "1g.10gb" 

23 assert specs["hunyuanframepackvae"].gpu == 1 

24 assert specs["podcasttranscript"].gpu == 0 

25 

26 

27def test_minimum_service_specs_reuse_container_resources_values() -> None: 

28 """CPU/memory/storage values should come from shared CONTAINER_RESOURCES.""" 

29 specs = get_minimum_service_container_specs(max_gpus=2) 

30 gemma = specs["gemma"] 

31 

32 assert gemma.cpu == CONTAINER_RESOURCES["gemma"].cpu 

33 assert gemma.memory_gib == CONTAINER_RESOURCES["gemma"].memory_gib 

34 assert gemma.ephemeral_storage_gib == CONTAINER_RESOURCES["gemma"].ephemeral_storage_gib 

35 

36 

37def test_model_container_resources_map_into_container_resources() -> None: 

38 """Model-keyed defaults should be exposed through string container names.""" 

39 for model, spec in MODEL_CONTAINER_RESOURCES.items(): 

40 container_name = MODEL_TO_CONTAINER_NAME[model] 

41 assert container_name in CONTAINER_RESOURCES 

42 assert CONTAINER_RESOURCES[container_name] == spec 

43 

44 

45def test_minimum_service_specs_gpu_scaling_edge_cases() -> None: 

46 """Scalable services should clamp GPU defaults and treat negative max_gpus as zero.""" 

47 specs_zero = get_minimum_service_container_specs(max_gpus=0) 

48 specs_one = get_minimum_service_container_specs(max_gpus=1) 

49 specs_negative = get_minimum_service_container_specs(max_gpus=-2) 

50 

51 assert specs_zero["gemma"].gpu == 0 

52 assert specs_one["gemma"].gpu == 1 

53 assert specs_negative["gemma"].gpu == 0 

54 # MIG containers always keep their MIG profile regardless of max_gpus. 

55 assert specs_zero["yolo"].gpu == "1g.10gb" 

56 assert specs_zero["kokoro"].gpu == "1g.10gb" 

57 assert specs_zero["realesrgan"].gpu == "1g.10gb"