add unit tests for katib sdk

Signed-off-by: Shashank Mittal <[email protected]>

Makefile

@@ -178,6 +178,7 @@ pytest: prepare-pytest prepare-pytest-testdata
 	PYTHONPATH=$(PYTHONPATH) pytest ./test/unit/v1beta1/suggestion --ignore=./test/unit/v1beta1/suggestion/test_skopt_service.py
 	PYTHONPATH=$(PYTHONPATH) pytest ./test/unit/v1beta1/earlystopping
 	PYTHONPATH=$(PYTHONPATH) pytest ./test/unit/v1beta1/metricscollector
+	PYTHONPATH=$(PYTHONPATH) pytest ./sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py
 
 # The skopt service doesn't work appropriately with Python 3.11.
 # So, we need to run the test with Python 3.9.

sdk/python/v1beta1/kubeflow/katib/api/katib_client.py

@@ -129,8 +129,9 @@ def create_experiment(
                 f"Failed to create Katib Experiment: {namespace}/{experiment_name}"
             )
 
-        # TODO (andreyvelich): Use proper logger.
-        print(f"Experiment {namespace}/{experiment_name} has been created")
+        success_message = f"Experiment {namespace}/{experiment_name} has been created"
+        print(success_message)
+        return success_message
 
         if self._is_ipython():
             if self.in_cluster:

sdk/python/v1beta1/kubeflow/katib/api/katib_client_test.py

@@ -0,0 +1,55 @@
+import multiprocessing
+import pytest
+from unittest.mock import patch, Mock
+from kubeflow.katib import models
+from kubeflow.katib.api.katib_client import KatibClient
+
+@pytest.fixture
+def katib_client():
+    return KatibClient()
+
+@pytest.fixture
+def valid_experiment():
+    return models.V1beta1Experiment(
+        metadata=models.V1ObjectMeta(name="test-experiment"),
+        spec=models.V1beta1ExperimentSpec()
+    )
+
+class TestCreateExperiment:
+    @patch('kubeflow.katib.api.katib_client.client.CustomObjectsApi.create_namespaced_custom_object')
+    def test_create_experiment_success(self, mock_create_namespaced_custom_object, katib_client, valid_experiment):
+        mock_create_namespaced_custom_object.return_value = {"metadata": {"name": "test-experiment"}}
+        experiment = katib_client.create_experiment(valid_experiment)
+        assert experiment == "Experiment default/test-experiment has been created"
+
+    @patch('kubeflow.katib.api.katib_client.client.CustomObjectsApi.create_namespaced_custom_object')
+    def test_create_experiment_timeout_error(self, mock_create_namespaced_custom_object, katib_client, valid_experiment):
+        mock_create_namespaced_custom_object.side_effect = multiprocessing.TimeoutError()
+        with pytest.raises(TimeoutError):
+            katib_client.create_experiment(valid_experiment, namespace="timeout")
+
+    @patch('kubeflow.katib.api.katib_client.client.CustomObjectsApi.create_namespaced_custom_object')
+    def test_create_experiment_runtime_error(self, mock_create_namespaced_custom_object, katib_client, valid_experiment):
+        mock_create_namespaced_custom_object.side_effect = RuntimeError()
+        with pytest.raises(RuntimeError):
+            katib_client.create_experiment(valid_experiment, namespace="runtime")
+
+class TestGetExperiment:
+    @patch('kubeflow.katib.api.katib_client.client.CustomObjectsApi.get_namespaced_custom_object')
+    def test_get_experiment_success(self, mock_get_namespaced_custom_object, katib_client, valid_experiment):
+        mock_response = {"metadata": {"name": "test-experiment"}, "spec": {}, "status": {}}
+        mock_get_namespaced_custom_object.return_value = mock_response
+        experiment = katib_client.get_experiment("test-experiment")
+        assert experiment == valid_experiment
+
+    @patch('kubeflow.katib.api.katib_client.client.CustomObjectsApi.get_namespaced_custom_object')
+    def test_get_experiment_timeout_error(self, mock_get_namespaced_custom_object, katib_client):
+        mock_get_namespaced_custom_object.side_effect = multiprocessing.TimeoutError()
+        with pytest.raises(TimeoutError):
+            katib_client.get_experiment("test-experiment", namespace="timeout")
+
+    @patch('kubeflow.katib.api.katib_client.client.CustomObjectsApi.get_namespaced_custom_object')
+    def test_get_experiment_runtime_error(self, mock_get_namespaced_custom_object, katib_client):
+        mock_get_namespaced_custom_object.side_effect = RuntimeError()
+        with pytest.raises(RuntimeError):
+            katib_client.get_experiment("test-experiment", namespace="runtime")