Coverage for tests/k8s_mock.py: 96%

1""" 

2A mock for Kubernetes modules used in testing. 

3""" 

4 

5from unittest.mock import MagicMock 

6from unittest.mock import AsyncMock 

7 

8from typing import Dict 

9from typing import Any 

10from typing import Tuple 

11 

12 

13class MockApiException(BaseException): 

14 """A mock for kubernetes_asyncio.client.exceptions.ApiException.""" 

15 def __init__( 

16 self, 

17 status: int, 

18 reason: str, 

19 body: Any = None, 

20 ) -> None: 

21 self.status = status 

22 self.reason = reason 

23 self.body = body 

24 super().__init__(f"ApiException: {status} {reason}") 

25 

26 

27class MockInvalidUrlClientError(BaseException): 

28 """A mock for kubernetes_asyncio.client.exceptions.InvalidUrlClientError.""" 

29 def __init__( 

30 self, 

31 message: str, 

32 ) -> None: 

33 self.message = message 

34 super().__init__(f"InvalidUrlClientError: {message}") 

35 

36 

37class K8sMock(MagicMock): 

38 """A mock for Kubernetes modules with specific mocked attributes and methods.""" 

39 

40 def __init__( 

41 self, 

42 *args: Tuple, 

43 **kwargs: Dict 

44 ) -> None: 

45 super().__init__(*args, **kwargs) 

46 

47 def get_sub_modules(self) -> Dict[str, Any]: 

48 # Async context manager for ApiClient 

49 mock_kubernetes_asyncio = MagicMock() 

50 mock_api_client = MagicMock() 

51 mock_api_client_cm = AsyncMock() 

52 mock_api_client_cm.__aenter__.return_value = mock_api_client 

53 mock_api_client_cm.__aexit__.return_value = None 

54 

55 mock_core_api = MagicMock() 

56 mock_core_api.list_namespaced_pod = AsyncMock(return_value=MagicMock(items=[])) 

57 mock_core_api.list_namespace = AsyncMock(return_value=MagicMock(items=[])) 

58 mock_core_api.list_node = AsyncMock(return_value=MagicMock(items=[])) 

59 mock_core_api.list_pod_for_all_namespaces = AsyncMock(return_value=MagicMock(items=[])) 

60 mock_core_api.list_service_for_all_namespaces = AsyncMock(return_value=MagicMock(items=[])) 

61 mock_core_api.delete_namespaced_pod = AsyncMock(return_value=MagicMock(status="Success")) 

62 mock_core_api.delete_namespaced_service = AsyncMock(return_value=MagicMock(status="Success")) 

63 mock_core_api.read_namespaced_service_account = AsyncMock() 

64 mock_core_api.read_namespaced_pod_log = AsyncMock(return_value="log line") 

65 mock_core_api.create_namespaced_pod = AsyncMock(return_value=[]) 

66 mock_core_api.create_namespaced_service = AsyncMock(return_value=MagicMock(status="Success")) 

67 

68 mock_custom_api = MagicMock() 

69 mock_custom_api.get_namespaced_custom_object = AsyncMock(return_value=MagicMock()) 

70 

71 mock_rbac_api = MagicMock() 

72 mock_rbac_api.read_cluster_role = AsyncMock() 

73 mock_rbac_api.read_cluster_role_binding = AsyncMock() 

74 

75 mock_k8s_client = MagicMock() 

76 mock_k8s_client.ApiClient.return_value = mock_api_client_cm 

77 mock_k8s_client.CoreV1Api.return_value = mock_core_api 

78 mock_k8s_client.CustomObjectsApi.return_value = mock_custom_api 

79 mock_k8s_client.RbacAuthorizationV1Api.return_value = mock_rbac_api 

80 

81 mock_k8s_config = MagicMock() 

82 mock_k8s_config.load_kube_config = AsyncMock() 

83 mock_k8s_config.load_incluster_config = MagicMock() 

84 mock_k8s_config.kube_config = MagicMock() 

85 

86 mock_exceptions = MagicMock() 

87 mock_exceptions.ApiException = MockApiException 

88 mock_exceptions.InvalidUrlClientError = MockInvalidUrlClientError 

89 

90 return { 

91 "kubernetes_asyncio": mock_kubernetes_asyncio, 

92 "kubernetes_asyncio.config": mock_k8s_config, 

93 "kubernetes_asyncio.client": mock_k8s_client, 

94 "kubernetes_asyncio.client.exceptions": mock_exceptions, 

95 "kubernetes_asyncio.config.config_exception": MagicMock(), 

96 }