Coverage for tests/streamwise_app/test_service_manager.py: 93%

1#!/usr/bin/env python3 

2""" 

3Unit tests for LMM Service Manager. 

4""" 

5import sys 

6import os 

7import pytest 

8 

9from http import HTTPStatus 

10 

11from unittest.mock import patch 

12from unittest.mock import MagicMock 

13from unittest.mock import AsyncMock 

14 

15# Add current path 

16sys.path.append(os.getcwd()) 

17 

18from tests.torch_mock import TorchMock 

19 

20from tests.test_utils import temp_sys_path 

21 

22mock_torch = TorchMock() 

23 

24mock_modules = {} 

25mock_modules.update(mock_torch.get_sub_modules()) 

26 

27with patch.dict(sys.modules, mock_modules): 

28 with temp_sys_path("apps"): 

29 from apps.client import LMMServiceManager 

30 from apps.client import ServiceNotFoundError 

31 

32 

33def _mock_session_http(service_name: str) -> MagicMock: 

34 """Mock the async aiohttp ClientSession for HTTP calls.""" 

35 response = MagicMock() 

36 response.status = HTTPStatus.OK 

37 response.json = AsyncMock(return_value={ 

38 service_name: { 

39 "status": "ok", 

40 "gpu": "A100", 

41 "world_size": 2, 

42 } 

43 }) 

44 cm = AsyncMock() 

45 cm.__aenter__.return_value = response 

46 cm.__aexit__.return_value = None 

47 session = MagicMock() 

48 session.get.return_value = cm 

49 session.close = AsyncMock() 

50 return session 

51 

52 

53@pytest.mark.asyncio 

54async def test_service_manager() -> None: 

55 service_manager = LMMServiceManager("streamwise") 

56 assert service_manager is not None 

57 

58 # Mocking responses from services 

59 service_manager.session = _mock_session_http("flux") 

60 

61 status, is_busy, gpu_model, num_gpus = await service_manager.get_container_status( 

62 service_name="flux", 

63 url="http://flux:8080/health" 

64 ) 

65 assert status == "ok" 

66 assert is_busy is False 

67 assert gpu_model == "A100" 

68 assert num_gpus == 2 

69 

70 await service_manager.start_updater() 

71 

72 service_manager.print_service_status() 

73 

74 with pytest.raises(ServiceNotFoundError): 

75 service_manager.get_service_url("fantasytalking") 

76 

77 with pytest.raises(ServiceNotFoundError): 

78 service_manager.get_service_urls("qwenimage") 

79 

80 await service_manager.stop() 

81 

82 

83def test_parse_arguments_https_args() -> None: 

84 """parse_arguments() accepts --certfile and --keyfile for HTTPS.""" 

85 with patch.dict(sys.modules, mock_torch.get_sub_modules()): 

86 with temp_sys_path("apps"): 

87 from apps.streamwise_app import StreamWiseApp 

88 from apps.streamwise_job import StreamWiseJob 

89 

90 class _TestApp(StreamWiseApp): 

91 def setup_routes(self) -> None: 

92 pass 

93 

94 def create_job(self, job_id: str, job_config: dict) -> StreamWiseJob: 

95 raise NotImplementedError 

96 

97 app_instance = _TestApp("test") 

98 with patch("sys.argv", ["app", "--certfile", "/tmp/cert.pem", "--keyfile", "/tmp/key.pem"]): 

99 args = app_instance.parse_arguments("test") 

100 assert args.certfile == "/tmp/cert.pem" 

101 assert args.keyfile == "/tmp/key.pem" 

102 

103 

104def test_parse_arguments_https_defaults() -> None: 

105 """parse_arguments() defaults certfile and keyfile to None (HTTP mode).""" 

106 with patch.dict(sys.modules, mock_torch.get_sub_modules()): 

107 with temp_sys_path("apps"): 

108 from apps.streamwise_app import StreamWiseApp 

109 from apps.streamwise_job import StreamWiseJob 

110 

111 class _TestApp(StreamWiseApp): 

112 def setup_routes(self) -> None: 

113 pass 

114 

115 def create_job(self, job_id: str, job_config: dict) -> StreamWiseJob: 

116 raise NotImplementedError 

117 

118 app_instance = _TestApp("test") 

119 with patch("sys.argv", ["app"]): 

120 args = app_instance.parse_arguments("test") 

121 assert args.certfile is None 

122 assert args.keyfile is None 

123 

124 

125@pytest.mark.asyncio 

126async def test_run_httpserver_https() -> None: 

127 """run_httpserver() configures Hypercorn SSL when certfile/keyfile are given.""" 

128 with patch.dict(sys.modules, mock_torch.get_sub_modules()): 

129 with temp_sys_path("apps"): 

130 from apps.streamwise_app import StreamWiseApp 

131 from apps.streamwise_job import StreamWiseJob 

132 

133 class _TestApp(StreamWiseApp): 

134 def setup_routes(self) -> None: 

135 pass 

136 

137 def create_job(self, job_id: str, job_config: dict) -> StreamWiseJob: 

138 raise NotImplementedError 

139 

140 app_instance = _TestApp("test") 

141 with patch("apps.streamwise_app.serve", new=AsyncMock()) as mock_serve: 

142 await app_instance.run_httpserver( 

143 host="127.0.0.1", 

144 port=9999, 

145 certfile="/tmp/cert.pem", 

146 keyfile="/tmp/key.pem", 

147 ) 

148 mock_serve.assert_awaited_once() 

149 config = mock_serve.call_args[0][1] 

150 assert config.certfile == "/tmp/cert.pem" 

151 assert config.keyfile == "/tmp/key.pem" 

152 

153 

154@pytest.mark.asyncio 

155async def test_run_httpserver_http() -> None: 

156 """run_httpserver() configures Hypercorn without SSL when no certfile is given.""" 

157 with patch.dict(sys.modules, mock_torch.get_sub_modules()): 

158 with temp_sys_path("apps"): 

159 from apps.streamwise_app import StreamWiseApp 

160 from apps.streamwise_job import StreamWiseJob 

161 

162 class _TestApp(StreamWiseApp): 

163 def setup_routes(self) -> None: 

164 pass 

165 

166 def create_job(self, job_id: str, job_config: dict) -> StreamWiseJob: 

167 raise NotImplementedError 

168 

169 app_instance = _TestApp("test") 

170 with patch("apps.streamwise_app.serve", new=AsyncMock()) as mock_serve: 

171 await app_instance.run_httpserver(host="127.0.0.1", port=9999) 

172 mock_serve.assert_awaited_once() 

173 config = mock_serve.call_args[0][1] 

174 assert getattr(config, "certfile", None) is None 

175 assert getattr(config, "keyfile", None) is None