LK: Added tests for get_castai_nodes and check_controller_replicas

tests/test_node_utils.py

@@ -2,13 +2,14 @@
 import unittest
 from unittest.mock import MagicMock
 from datetime import datetime, timedelta, timezone
-from kubernetes.client import V1Node, V1ObjectMeta, V1Pod, V1PodList
+from kubernetes.client import V1Node, V1ObjectMeta, V1Pod, V1PodList, V1NodeList
 
 # Add src directory to the system path
 sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../src')))
 
 from node_utils import is_node_older_than
 from node_utils import is_node_running_critical_pods
+from node_utils import get_cast_ai_nodes
 
 # Mock the CRITICAL_WORKLOADS to match the tests
 CRITICAL_WORKLOADS = ["app.kubernetes.io/name=castai-agent", "app.kubernetes.io/name=castai-cluster-controller"]
@@ -87,5 +88,45 @@ def test_is_node_running_critical_pods_empty(self):
         result = is_node_running_critical_pods(v1, node_name)
         self.assertFalse(result, "Node with no pods should return False")
 
+    def test_get_cast_ai_nodes_positive(self):
+        # Mock the CoreV1Api instance and its list_node method
+        v1 = MagicMock()
+
+        # Create a list of nodes that match the CAST AI managed label
+        node1 = V1Node(metadata=V1ObjectMeta(name="node1"))
+        node2 = V1Node(metadata=V1ObjectMeta(name="node2"))
+        node_list = V1NodeList(items=[node1, node2])
+
+        v1.list_node.return_value = node_list
+
+        # Test the function
+        result = get_cast_ai_nodes(v1)
+
+        # Verify the label selector was used
+        v1.list_node.assert_called_with(label_selector="provisioner.cast.ai/managed-by=cast.ai")
+
+        # Verify the result
+        self.assertEqual(len(result), 2, "Should return 2 nodes")
+        self.assertEqual(result[0].metadata.name, "node1", "First node name should be 'node1'")
+        self.assertEqual(result[1].metadata.name, "node2", "Second node name should be 'node2'")
+
+    def test_get_cast_ai_nodes_negative(self):
+        # Mock the CoreV1Api instance and its list_node method
+        v1 = MagicMock()
+
+        # Create an empty list of nodes
+        node_list = V1NodeList(items=[])
+
+        v1.list_node.return_value = node_list
+
+        # Test the function
+        result = get_cast_ai_nodes(v1)
+
+        # Verify the label selector was used
+        v1.list_node.assert_called_with(label_selector="provisioner.cast.ai/managed-by=cast.ai")
+
+        # Verify the result
+        self.assertEqual(len(result), 0, "Should return 0 nodes")
+
 if __name__ == '__main__':
     unittest.main()

tests/test_pod_utils.py

@@ -0,0 +1,79 @@
+import unittest
+from unittest.mock import MagicMock
+from typing import List, Optional, Tuple
+from kubernetes.client import CoreV1Api, V1Pod, V1PodList, V1ObjectMeta, V1OwnerReference
+import sys
+import os
+
+# Add src directory to the system path
+sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../src')))
+
+from pod_utils import check_controller_replicas
+
+class TestPodUtils(unittest.TestCase):
+
+    def test_check_controller_replicas_happy_path(self):
+        v1 = MagicMock()
+        node_name = "test-node"
+        namespace = "default"
+
+        # Create pods that match the controller replicas
+        owner_reference = V1OwnerReference(api_version="v1", kind="ReplicaSet", name="test-controller", uid="123")
+        pod1 = V1Pod(metadata=V1ObjectMeta(name="pod1", namespace=namespace, owner_references=[owner_reference], labels={"app": "test-app"}))
+        pod2 = V1Pod(metadata=V1ObjectMeta(name="pod2", namespace=namespace, owner_references=[owner_reference], labels={"app": "test-app"}))
+        pod_list = V1PodList(items=[pod1, pod2])
+
+        v1.list_pod_for_all_namespaces.return_value = pod_list
+        v1.list_namespaced_pod.return_value = pod_list
+
+        result = check_controller_replicas(v1, node_name)
+        self.assertEqual(result, ("ReplicaSet", "test-controller", namespace, [pod1, pod2]))
+
+    def test_check_controller_replicas_no_pods(self):
+        v1 = MagicMock()
+        node_name = "test-node"
+
+        # Create an empty list of pods
+        pod_list = V1PodList(items=[])
+
+        v1.list_pod_for_all_namespaces.return_value = pod_list
+
+        result = check_controller_replicas(v1, node_name)
+        self.assertEqual(result, (None, None, None, None))
+
+    def test_check_controller_replicas_single_replica(self):
+        v1 = MagicMock()
+        node_name = "test-node"
+        namespace = "default"
+
+        # Create a pod with a single replica
+        owner_reference = V1OwnerReference(api_version="v1", kind="ReplicaSet", name="test-controller", uid="123")
+        pod = V1Pod(metadata=V1ObjectMeta(name="pod1", namespace=namespace, owner_references=[owner_reference], labels={"app": "test-app"}))
+        pod_list = V1PodList(items=[pod])
+
+        v1.list_pod_for_all_namespaces.return_value = pod_list
+
+        result = check_controller_replicas(v1, node_name)
+        self.assertEqual(result, (None, None, None, None))
+
+    def test_check_controller_replicas_not_all_replicas_on_node(self):
+        v1 = MagicMock()
+        node_name = "test-node"
+        namespace = "default"
+
+        # Create pods that match the controller replicas
+        owner_reference = V1OwnerReference(api_version="v1", kind="ReplicaSet", name="test-controller", uid="123")
+        pod1 = V1Pod(metadata=V1ObjectMeta(name="pod1", namespace=namespace, owner_references=[owner_reference], labels={"app": "test-app"}))
+        pod2 = V1Pod(metadata=V1ObjectMeta(name="pod2", namespace=namespace, owner_references=[owner_reference], labels={"app": "test-app"}))
+        pod3 = V1Pod(metadata=V1ObjectMeta(name="pod3", namespace=namespace, owner_references=[owner_reference], labels={"app": "test-app"}))
+        pod_list_node = V1PodList(items=[pod1, pod2])
+        pod_list_all = V1PodList(items=[pod1, pod2, pod3])
+
+        v1.list_pod_for_all_namespaces.return_value = pod_list_node
+        v1.list_namespaced_pod.return_value = pod_list_all
+
+        result = check_controller_replicas(v1, node_name)
+        self.assertEqual(result, (None, None, None, None))
+
+if __name__ == '__main__':
+    unittest.main()