Coverage for tests/test_k8s_utils.py: 96%

1""" 

2Test K8s utilities. 

3""" 

4 

5import logging 

6import pytest 

7 

8from unittest.mock import MagicMock 

9 

10from kubernetes_asyncio.config.config_exception import ConfigException 

11from aiohttp.client_exceptions import InvalidUrlClientError 

12 

13from k8s_utils import NoRunnableContainerError 

14from k8s_utils import parse_k8s_resource_quantity 

15from k8s_utils import parse_mig_resources 

16from k8s_utils import get_k8s_nodes 

17from k8s_utils import get_k8s_pods 

18from k8s_utils import get_k8s_services 

19from k8s_utils import get_svc_name_from_container_name 

20from k8s_utils import get_k8s_services_ns 

21from k8s_utils import K8sContainer 

22from k8s_utils import K8sService 

23from k8s_utils import NoActiveContainerError 

24from k8s_utils import is_k8s_node_ready 

25import k8s_utils 

26 

27 

28def test_parse_k8s_resource_quantity() -> None: 

29 assert parse_k8s_resource_quantity("1000m") == 1 

30 assert parse_k8s_resource_quantity("200m") == 0.2 

31 assert parse_k8s_resource_quantity("1") == 1 

32 assert parse_k8s_resource_quantity("2") == 2 

33 assert parse_k8s_resource_quantity("200Ki") == 200 * 1024 

34 assert parse_k8s_resource_quantity("512Mi") == 512 * (1024 ** 2) 

35 assert parse_k8s_resource_quantity("1Gi") == 1024 ** 3 

36 assert parse_k8s_resource_quantity("4Gi") == 4 * (1024 ** 3) 

37 assert parse_k8s_resource_quantity("0.2") == 0.2 

38 assert parse_k8s_resource_quantity("bogus") == 0 

39 

40 

41def test_parse_mig_resources() -> None: 

42 assert parse_mig_resources(None) == {} 

43 

44 resources = {"cpu": "4", "nvidia.com/gpu": "2"} 

45 assert parse_mig_resources(resources) == {} 

46 

47 resources = { 

48 "cpu": "4", 

49 "memory": "8Gi", 

50 "nvidia.com/gpu": "0", 

51 "nvidia.com/mig-1g.5gb": "7", 

52 "nvidia.com/mig-2g.10gb": "3", 

53 } 

54 assert parse_mig_resources(resources) == {"1g.5gb": 7, "2g.10gb": 3} 

55 

56 resources = { 

57 "cpu": "4", 

58 "memory": "8Gi", 

59 "nvidia.com/gpu": "7", 

60 "nvidia.com/mig-1g.10gb": "3", 

61 "nvidia.com/mig-2g.20gb": "2", 

62 } 

63 assert parse_mig_resources(resources) == {"1g.10gb": 3, "2g.20gb": 2} 

64 

65 

66@pytest.mark.asyncio 

67async def test_get_k8s_nodes() -> None: 

68 try: 

69 nodes = await get_k8s_nodes() 

70 assert nodes is not None 

71 # assert len(nodes) == 0 

72 except ConfigException: 

73 logging.info("No K8s setup") 

74 

75 

76@pytest.mark.asyncio 

77async def test_get_k8s_pods() -> None: 

78 try: 

79 pods = await get_k8s_pods() 

80 assert pods is not None 

81 # assert len(pods) == 0 

82 except ConfigException: 

83 logging.info("No K8s setup") 

84 

85 try: 

86 pods = await get_k8s_pods(context_name="unittest") 

87 assert pods is not None 

88 except InvalidUrlClientError: 

89 logging.info("No K8s setup for unittest context") 

90 

91 

92@pytest.mark.asyncio 

93async def test_get_k8s_services() -> None: 

94 try: 

95 services = await get_k8s_services() 

96 assert services is not None 

97 # assert len(services) == 0 

98 except ConfigException: 

99 logging.info("No K8s setup") 

100 

101 try: 

102 services = await get_k8s_services(context_name="unittest") 

103 assert services is not None 

104 except InvalidUrlClientError: 

105 logging.info("No K8s setup for unittest context") 

106 

107 

108def test_data_model() -> None: 

109 # Container 

110 container = K8sContainer("flux_1234") 

111 assert container.name == "flux_1234" 

112 container.set_gpu_model("A100") 

113 assert container.get_gpu_model() == "A100" 

114 container.set_num_gpus(2) 

115 assert container.get_num_gpus() == 2 

116 assert container.is_active() is False 

117 assert container.is_busy() is False 

118 container.busy = True 

119 assert container.is_busy() is True 

120 container.busy = False 

121 

122 container.status = "ok" 

123 container.ip = "1.2.3.4" 

124 container.port = 8080 

125 assert container.get_url() == "http://1.2.3.4:8080" 

126 

127 assert container.is_active() is True 

128 assert container.is_busy() is False 

129 assert str(container) == "K8sContainer(name=flux_1234, status=ok, busy=False, ip=1.2.3.4, port=8080, resources={})" 

130 logging.info(f"Container: {container}") 

131 

132 # Service 

133 service = K8sService("flux") 

134 with pytest.raises(NoActiveContainerError): 

135 service.get_active_containers() 

136 with pytest.raises(NoActiveContainerError): 

137 service.get_runnable_containers() 

138 

139 service.add_container(container) 

140 active_containers = service.get_active_containers() 

141 assert len(active_containers) == 1 

142 

143 runnable_containers = service.get_runnable_containers() 

144 assert len(runnable_containers) == 1 

145 service.containers[0].busy = True 

146 with pytest.raises(NoRunnableContainerError, match="flux"): 

147 service.get_runnable_containers() 

148 service.containers[0].busy = False 

149 runnable_containers = service.get_runnable_containers() 

150 assert len(runnable_containers) == 1 

151 

152 best_containers = service.get_best_containers() 

153 assert best_containers is not None 

154 assert len(best_containers) == 1 

155 

156 best_containers = service.get_best_containers(exclude_busy=False) 

157 assert best_containers is not None 

158 assert len(best_containers) == 1 

159 

160 best_containers = service.get_best_containers(excluded_containers=[container]) 

161 assert best_containers is not None 

162 assert best_containers == [] 

163 

164 best_container = service.get_best_container() 

165 assert best_container is not None 

166 

167 assert str(service) == "K8sService(name=flux, containers=1)" 

168 

169 container.status = None 

170 with pytest.raises(NoActiveContainerError): 

171 service.get_active_containers() 

172 

173 

174@pytest.mark.asyncio 

175async def test_get_k8s_services_ns() -> None: 

176 try: 

177 services = await get_k8s_services_ns() 

178 assert services is not None 

179 except ConfigException: 

180 logging.info("No K8s setup") 

181 

182 try: 

183 services = await get_k8s_services_ns("unittest") 

184 assert services is not None 

185 except InvalidUrlClientError: 

186 logging.info("No K8s setup for unittest context") 

187 

188 

189def test_is_k8s_node_ready() -> None: 

190 node_mock = MagicMock() 

191 assert is_k8s_node_ready(node_mock) is False 

192 

193 node_mock.status.conditions = [ 

194 MagicMock(type="Ready", status="True"), 

195 ] 

196 assert is_k8s_node_ready(node_mock) is True 

197 

198 

199def test_get_svc_name_from_container_name() -> None: 

200 assert get_svc_name_from_container_name("flux-1234") == "flux_1234" 

201 assert get_svc_name_from_container_name("model-serving-abcde_5678") == "model_serving_abcde_5678" 

202 

203 

204def test_service_scheme_default() -> None: 

205 """SERVICE_SCHEME defaults to http.""" 

206 assert k8s_utils.SERVICE_SCHEME == "http" 

207 

208 

209def test_set_service_scheme_https() -> None: 

210 """set_service_scheme switches K8sContainer.get_url() to https.""" 

211 original = k8s_utils.SERVICE_SCHEME 

212 try: 

213 k8s_utils.set_service_scheme("https") 

214 assert k8s_utils.SERVICE_SCHEME == "https" 

215 container = K8sContainer("test") 

216 container.ip = "1.2.3.4" 

217 container.port = 8080 

218 assert container.get_url() == "https://1.2.3.4:8080" 

219 finally: 

220 k8s_utils.set_service_scheme(original) 

221 

222 

223def test_set_service_scheme_http() -> None: 

224 """set_service_scheme can be reset back to http.""" 

225 original = k8s_utils.SERVICE_SCHEME 

226 try: 

227 k8s_utils.set_service_scheme("https") 

228 k8s_utils.set_service_scheme("http") 

229 assert k8s_utils.SERVICE_SCHEME == "http" 

230 container = K8sContainer("test") 

231 container.ip = "1.2.3.4" 

232 container.port = 8080 

233 assert container.get_url() == "http://1.2.3.4:8080" 

234 finally: 

235 k8s_utils.set_service_scheme(original) 

236 

237 

238def test_set_service_scheme_invalid() -> None: 

239 """set_service_scheme raises ValueError for unsupported schemes.""" 

240 with pytest.raises(ValueError, match="Invalid service scheme"): 

241 k8s_utils.set_service_scheme("ftp") 

242 

243 

244def test_verify_ssl_default() -> None: 

245 """VERIFY_SSL defaults to True.""" 

246 assert k8s_utils.VERIFY_SSL is True 

247 

248 

249def test_set_verify_ssl_false() -> None: 

250 """set_verify_ssl(False) disables SSL certificate verification.""" 

251 original = k8s_utils.VERIFY_SSL 

252 try: 

253 k8s_utils.set_verify_ssl(False) 

254 assert k8s_utils.VERIFY_SSL is False 

255 finally: 

256 k8s_utils.set_verify_ssl(original) 

257 

258 

259def test_set_verify_ssl_true() -> None: 

260 """set_verify_ssl(True) re-enables SSL certificate verification.""" 

261 original = k8s_utils.VERIFY_SSL 

262 try: 

263 k8s_utils.set_verify_ssl(False) 

264 k8s_utils.set_verify_ssl(True) 

265 assert k8s_utils.VERIFY_SSL is True 

266 finally: 

267 k8s_utils.set_verify_ssl(original)