[SDK] Unit tests for TrainingClient APIs - get_job_pod_names and update_job

sdk/python/kubeflow/training/api/training_client_test.py

@@ -17,11 +17,13 @@
 from kubernetes.client import V1ResourceRequirements
 import pytest
 
-LIST_RESPONSE = [{"metadata": {"name": "Dummy V1PodList"}}]
+LIST_RESPONSE = [
+    {"metadata": {"name": "Dummy V1PodList"}},
+]
 TEST_NAME = "test"
 
 
-def create_namespaced_custom_object_response(*args, **kwargs):
+def conditional_error_handler(*args, **kwargs):
     if args[2] == "timeout":
         raise multiprocessing.TimeoutError()
     elif args[2] == "runtime":
@@ -31,12 +33,28 @@ def create_namespaced_custom_object_response(*args, **kwargs):
 def list_namespaced_pod_response(*args, **kwargs):
     class MockResponse:
         def get(self, timeout):
-            # Simulate a response from the Kubernetes API, and pass timeout for verification
+            """
+            Simulates Kubernetes API response for listing namespaced pods,
+            and pass timeout for verification
+
+            :return:
+                - If `args[0] == "timeout"`, raises `TimeoutError`.
+                - If `args[0] == "runtime"`, raises `Exception`.
+                - If `args[0] == "mock_pod_obj"`, returns a mock pod object with `metadata.name = "Dummy V1PodList"`.
+                - If `args[0] == "no_pods"`, returns an empty list of pods.
+                - Otherwise, returns a default list of dicts representing pods, with `timeout` included, for testing.
+            """
             LIST_RESPONSE[0]["timeout"] = timeout
             if args[0] == "timeout":
                 raise multiprocessing.TimeoutError()
             if args[0] == "runtime":
                 raise Exception()
+            if args[0] == "mock_pod_obj":
+                pod_obj = Mock(metadata=Mock())
+                pod_obj.metadata.name = "Dummy V1PodList"
+                return Mock(items=[pod_obj])
+            if args[0] == "no_pods":
+                return Mock(items=[])
             return Mock(items=LIST_RESPONSE)
 
     return MockResponse()
@@ -250,14 +268,72 @@ def __init__(self, kind) -> None:
 ]
 
 
+test_data_get_job_pod_names = [
+    (
+        "valid flow",
+        {
+            "name": TEST_NAME,
+            "namespace": "mock_pod_obj",
+        },
+        ["Dummy V1PodList"],
+    ),
+    (
+        "valid flow with no pods available",
+        {
+            "name": TEST_NAME,
+            "namespace": "no_pods",
+        },
+        [],
+    ),
+]
+
+
+test_data_update_job = [
+    (
+        "valid flow",
+        {
+            "name": TEST_NAME,
+            "job": create_job(),
+        },
+        "No output",
+    ),
+    (
+        "invalid job_kind",
+        {
+            "name": TEST_NAME,
+            "job": create_job(),
+            "job_kind": "invalid_job_kind",
+        },
+        ValueError,
+    ),
+    (
+        "invalid flow with TimeoutError",
+        {
+            "name": TEST_NAME,
+            "namespace": "timeout",
+            "job": create_job(),
+        },
+        TimeoutError,
+    ),
+    (
+        "invalid flow with RuntimeError",
+        {
+            "name": TEST_NAME,
+            "namespace": "runtime",
+            "job": create_job(),
+        },
+        RuntimeError,
+    ),
+]
+
+
 @pytest.fixture
 def training_client():
     with patch(
         "kubernetes.client.CustomObjectsApi",
         return_value=Mock(
-            create_namespaced_custom_object=Mock(
-                side_effect=create_namespaced_custom_object_response
-            )
+            create_namespaced_custom_object=Mock(side_effect=conditional_error_handler),
+            patch_namespaced_custom_object=Mock(side_effect=conditional_error_handler),
         ),
     ), patch(
         "kubernetes.client.CoreV1Api",
@@ -309,3 +385,40 @@ def test_get_job_pods(
     except Exception as e:
         assert type(e) is expected_output
     print("test execution complete")
+
+
+@pytest.mark.parametrize(
+    "test_name,kwargs,expected_output",
+    test_data_get_job_pod_names,
+)
+def test_get_job_pod_names(training_client, test_name, kwargs, expected_output):
+    """
+    test get_job_pod_names function of training client
+    """
+    print("Executing test:", test_name)
+    out = training_client.get_job_pod_names(**kwargs)
+    assert out == expected_output
+    print("test execution complete")
+
+
+@pytest.mark.parametrize("test_name,kwargs,expected_output", test_data_update_job)
+def test_update_job(training_client, test_name, kwargs, expected_output):
+    """
+    test update_job function of training client
+    """
+    print("Executing test:", test_name)
+    try:
+        training_client.update_job(**kwargs)
+        training_client.custom_api.patch_namespaced_custom_object.assert_called_with(
+            constants.GROUP,
+            constants.VERSION,
+            kwargs.get("namespace", constants.DEFAULT_NAMESPACE),
+            constants.JOB_PARAMETERS[kwargs.get("job_kind", training_client.job_kind)][
+                "plural"
+            ],
+            kwargs.get("name"),
+            kwargs.get("job"),
+        )
+    except Exception as e:
+        assert type(e) is expected_output
+    print("test execution complete")