chore: addressed comments

Author: sai-sankalp-reddy

src/hyperpod_cli/commands/job.py

@@ -199,27 +199,9 @@ def list_jobs(
     try:
         logger.debug("Listing training jobs")
         result = list_training_job_service.list_training_jobs(
-            namespace, all_namespaces, selector
+            namespace, all_namespaces, selector, output
         )
-        result_dict = json.loads(result)
-
-        jobs = []
-        if "jobs" in result_dict and isinstance(result_dict["jobs"], list):
-            jobs = [job.values() for job in result_dict["jobs"]]
-
-        headers = [
-            "Name",
-            "Namespace",
-            "CreationTime",
-            "State"
-        ]
-
-        if output == OutputFormat.TABLE.value:
-            click.echo(tabulate(jobs, headers=headers, tablefmt="presto"))
-        elif output == OutputFormat.JSON.value:
-            json_list = [dict(zip(headers, value)) for value in jobs]
-            json_jobs = {"jobs": json_list}
-            click.echo(json.dumps(json_jobs, indent=4))
+        click.echo(result)
     except Exception as e:
         sys.exit(f"Unexpected error happens when trying to list training job : {e}")
 

src/hyperpod_cli/service/list_training_jobs.py

@@ -10,6 +10,7 @@
 # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
+import copy
 from typing import List, Optional
 
 import json
@@ -23,9 +24,10 @@
     V1ResourceAttributes
 )
 
-from hyperpod_cli.constants.command_constants import KUEUE_WORKLOAD_PRIORITY_CLASS_LABEL_KEY
+from hyperpod_cli.constants.command_constants import KUEUE_WORKLOAD_PRIORITY_CLASS_LABEL_KEY, OutputFormat
 from hyperpod_cli.constants.pytorch_constants import PYTORCH_CUSTOM_OBJECT_GROUP, PYTORCH_CUSTOM_OBJECT_PLURAL
 from hyperpod_cli.service.discover_namespaces import DiscoverNamespaces
+from tabulate import tabulate
 
 
 class ListTrainingJobs:
@@ -37,14 +39,14 @@ def list_training_jobs(
         namespace: Optional[str],
         all_namespaces: Optional[bool],
         selector: Optional[str],
+        output: Optional[str],
     ) -> str:
         """
         List training job provided by the user in the specified namespace.
         If namespace is not provided job are listed the default namespace in user context
         If all_namespace is true we will list training job from all namespaces that user has access
         Selector when specified will filter list of job addisional based on labels filter provided
         """
-
         k8s_client = KubernetesClient()
 
         jobs: List = []
@@ -80,10 +82,12 @@ def list_training_jobs(
         except ApiException as e:
             raise RuntimeError(f"Unexpected API error: {e.reason} ({e.status})")
 
-        return self._generate_list_training_job_output(jobs)
+        return self._generate_list_training_job_output(jobs, output)
 
-    def _generate_list_training_job_output(self, jobs: List):
+    def _generate_list_training_job_output(self, jobs: List, output: Optional[str]):
         output_jobs = {"jobs": []}
+        priority_header_required = False
+
         for job in jobs:
             if job.get("metadata"):
                 name = job.get("metadata").get("name")
@@ -104,10 +108,34 @@ def _generate_list_training_job_output(self, jobs: List):
 
                 if priority is not None:
                     job_summary["priority"] = priority
+                    priority_header_required = True
 
                 output_jobs["jobs"].append(job_summary)
 
-        return json.dumps(output_jobs, indent=1, sort_keys=False)
+        if output == OutputFormat.TABLE.value:
+            return self._generate_table(output_jobs, priority_header_required)
+        return json.dumps(output_jobs, indent=4, sort_keys=False)
+
+    def _generate_table(self, output_jobs, priority_header_required):
+        headers = [
+                "Name",
+                "Namespace",
+                "CreationTime",
+                "State"
+            ]
+
+        if priority_header_required:
+            headers.append("Priority")
+
+        jobs = []
+        if "jobs" in output_jobs and isinstance(output_jobs["jobs"], list):
+            for job in output_jobs["jobs"]:
+                job_values = list(job.values())
+                if priority_header_required and len(job_values) == 4:
+                    job_values.append("NA")
+                jobs.append(job_values)
+
+        return tabulate(jobs, headers=headers, tablefmt="presto")
 
     def _get_job_status(self, status: List) -> Optional[str]:
         current_status = None

test/unit_tests/service/test_list_training_jobs_service.py

@@ -23,6 +23,7 @@
 )
 
 from kubernetes.client.rest import ApiException
+from tabulate import tabulate
 
 SAMPLE_OUTPUT = {
     "items": [
@@ -152,7 +153,7 @@ def test_list_training_jobs_with_namespace(
         mock_kubernetes_client.return_value = self.mock_k8s_client
         self.mock_k8s_client.list_training_jobs.return_value = SAMPLE_OUTPUT
         result = self.mock_list_training_jobs.list_training_jobs(
-            "namespace", None, None
+            "namespace", None, None, None
         )
         self.assertIn("test-name", result)
         self.assertIn("test-name1", result)
@@ -166,7 +167,7 @@ def test_list_training_jobs_without_namespace(
         mock_kubernetes_client.return_value = self.mock_k8s_client
         self.mock_k8s_client.get_current_context_namespace.return_value = "namespace"
         self.mock_k8s_client.list_training_jobs.return_value = SAMPLE_OUTPUT
-        result = self.mock_list_training_jobs.list_training_jobs(None, None, None)
+        result = self.mock_list_training_jobs.list_training_jobs(None, None, None, None)
         self.assertIn("test-name", result)
         self.assertIn("test-name1", result)
         self.assertIn("Running", result)
@@ -182,7 +183,7 @@ def test_list_training_jobs_with_namespace_auto_discover(
         mock_discover_accessible_namespace.return_value = "discovered-namespace"
         self.mock_k8s_client.get_current_context_namespace.return_value = None
         self.mock_k8s_client.list_training_jobs.return_value = SAMPLE_OUTPUT
-        result = self.mock_list_training_jobs.list_training_jobs(None, None, None)
+        result = self.mock_list_training_jobs.list_training_jobs(None, None, None, None)
         mock_discover_accessible_namespace.assert_called_once_with(
             V1ResourceAttributes(
                 verb="list",
@@ -209,7 +210,7 @@ def test_list_training_jobs_without_namespace_no_jobs(
         mock_kubernetes_client.return_value = self.mock_k8s_client
         self.mock_k8s_client.get_current_context_namespace.return_value = "namespace"
         self.mock_k8s_client.list_training_jobs.return_value = {"items": []}
-        result = self.mock_list_training_jobs.list_training_jobs(None, None, None)
+        result = self.mock_list_training_jobs.list_training_jobs(None, None, None, None)
         self.assertNotIn("test-name", result)
         self.assertNotIn("test-name1", result)
 
@@ -221,7 +222,7 @@ def test_list_training_jobs_all_namespace(
         mock_kubernetes_client.return_value = self.mock_k8s_client
         self.mock_k8s_client.list_namespaces.return_value = ["namespace"]
         self.mock_k8s_client.list_training_jobs.return_value = SAMPLE_OUTPUT
-        result = self.mock_list_training_jobs.list_training_jobs(None, True, None)
+        result = self.mock_list_training_jobs.list_training_jobs(None, True, None, None)
         self.assertIn("test-name", result)
         self.assertIn("test-name1", result)
 
@@ -236,7 +237,7 @@ def test_list_training_jobs_all_namespace_api_exception(
             status="Failed", reason="unexpected"
         )
         with self.assertRaises(RuntimeError):
-            self.mock_list_training_jobs.list_training_jobs(None, True, None)
+            self.mock_list_training_jobs.list_training_jobs(None, True, None, None)
 
     @mock.patch("hyperpod_cli.clients.kubernetes_client.KubernetesClient.__new__")
     def test_list_training_jobs_all_namespace_no_jobs(
@@ -246,7 +247,7 @@ def test_list_training_jobs_all_namespace_no_jobs(
         mock_kubernetes_client.return_value = self.mock_k8s_client
         self.mock_k8s_client.list_namespaces.return_value = ["namespace"]
         self.mock_k8s_client.list_training_jobs.return_value = {"items": []}
-        result = self.mock_list_training_jobs.list_training_jobs(None, True, None)
+        result = self.mock_list_training_jobs.list_training_jobs(None, True, None, None)
         self.assertNotIn("test-name", result)
         self.assertNotIn("test-name1", result)
 
@@ -258,7 +259,7 @@ def test_list_training_jobs_all_namespace_missing_metadata(
         mock_kubernetes_client.return_value = self.mock_k8s_client
         self.mock_k8s_client.list_namespaces.return_value = ["namespace"]
         self.mock_k8s_client.list_training_jobs.return_value = INVALID_OUTPUT
-        result = self.mock_list_training_jobs.list_training_jobs(None, True, None)
+        result = self.mock_list_training_jobs.list_training_jobs(None, True, None, None)
         self.assertNotIn("name", result)
 
     @mock.patch("hyperpod_cli.clients.kubernetes_client.KubernetesClient.__new__")
@@ -269,5 +270,94 @@ def test_list_training_jobs_all_namespace_missing_status(
         mock_kubernetes_client.return_value = self.mock_k8s_client
         self.mock_k8s_client.list_namespaces.return_value = ["namespace"]
         self.mock_k8s_client.list_training_jobs.return_value = OUTPUT_WITHOUT_STATUS
-        result = self.mock_list_training_jobs.list_training_jobs(None, True, None)
+        result = self.mock_list_training_jobs.list_training_jobs(None, True, None, None)
         self.assertNotIn("State: null", result)
+
+    def test_generate_table_with_no_priority_header_and_values(self):
+        list_training_jobs = ListTrainingJobs()
+        output_jobs = {
+            "jobs": [
+                {
+                    "Name": "job1",
+                    "Namespace": "namespace1",
+                    "CreationTime": "2023-01-01T00:00:00Z",
+                    "State": "Running"
+                }
+            ]
+        }
+        priority_header_required = False
+
+        result = list_training_jobs._generate_table(output_jobs, priority_header_required)
+
+        expected_headers = ["Name", "Namespace", "CreationTime", "State"]
+        expected_jobs = [["job1", "namespace1", "2023-01-01T00:00:00Z", "Running"]]
+        expected_result = tabulate(expected_jobs, headers=expected_headers, tablefmt="presto")
+
+        assert result == expected_result
+
+    def test_generate_table_with_priority_header_and_priority_values(self):
+        list_training_jobs = ListTrainingJobs()
+        output_jobs = {
+            "jobs": [
+                {
+                    "Name": "job1",
+                    "Namespace": "namespace1",
+                    "CreationTime": "2023-01-01T00:00:00Z",
+                    "State": "Running",
+                    "priority": "high"
+                },
+                {
+                    "Name": "job2",
+                    "Namespace": "namespace2",
+                    "CreationTime": "2023-01-02T00:00:00Z",
+                    "State": "Completed",
+                    "priority": "low"
+                }
+            ]
+        }
+        priority_header_required = True
+
+        result = list_training_jobs._generate_table(output_jobs, priority_header_required)
+
+        expected_headers = ["Name", "Namespace", "CreationTime", "State", "Priority"]
+        expected_jobs = [
+            ["job1", "namespace1", "2023-01-01T00:00:00Z", "Running", "high"],
+            ["job2", "namespace2", "2023-01-02T00:00:00Z", "Completed", "low"]
+        ]
+        expected_result = tabulate(expected_jobs, headers=expected_headers, tablefmt="presto")
+
+        assert result == expected_result
+
+    def test_generate_table_with_priority_header_but_no_priority_value(self):
+        list_training_jobs = ListTrainingJobs()
+        output_jobs = {
+            "jobs": [
+                {
+                    "Name": "job1",
+                    "Namespace": "namespace1",
+                    "CreationTime": "2023-01-01T00:00:00Z",
+                    "State": "Running"
+                }
+            ]
+        }
+        priority_header_required = True
+
+        result = list_training_jobs._generate_table(output_jobs, priority_header_required)
+
+        expected_headers = ["Name", "Namespace", "CreationTime", "State", "Priority"]
+        expected_jobs = [["job1", "namespace1", "2023-01-01T00:00:00Z", "Running", "NA"]]
+        expected_result = tabulate(expected_jobs, headers=expected_headers, tablefmt="presto")
+
+        assert result == expected_result
+
+    def test_generate_table_empty_jobs(self):
+        list_training_jobs = ListTrainingJobs()
+        output_jobs = {"jobs": []}
+        priority_header_required = False
+
+        result = list_training_jobs._generate_table(output_jobs, priority_header_required)
+
+        expected_headers = ["Name", "Namespace", "CreationTime", "State"]
+        expected_result = tabulate([], headers=expected_headers, tablefmt="presto")
+
+        assert result == expected_result

test/unit_tests/test_job.py

@@ -10,7 +10,6 @@
 # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
-import json
 import unittest
 import click
 import subprocess
@@ -160,7 +159,7 @@ def test_list_job_happy_case(
         mock_list_training_job_service: mock.Mock,
     ):
         mock_list_training_job_service.return_value = self.mock_list_jobs
-        mock_list_training_job_service_and_list_jobs.return_value = json.dumps({"jobs": []})
+        mock_list_training_job_service_and_list_jobs.return_value = {"jobs": []}
         result = self.runner.invoke(list_jobs)
         self.assertEqual(result.exit_code, 0)
         self.assertIn("jobs", result.output)
@@ -177,60 +176,12 @@ def test_list_job_happy_case_debug_mode(
         mock_list_training_job_service: mock.Mock,
     ):
         mock_list_training_job_service.return_value = self.mock_list_jobs
-        mock_list_training_job_service_and_list_jobs.return_value = json.dumps({"jobs": []})
+        mock_list_training_job_service_and_list_jobs.return_value = {"jobs": []}
         result = self.runner.invoke(list_jobs, ["--debug"])
         self.assertEqual(result.exit_code, 0)
         self.assertIn("jobs", result.output)
         mock_debug.assert_called()
 
-    @mock.patch("hyperpod_cli.service.list_training_jobs.ListTrainingJobs")
-    @mock.patch(
-        "hyperpod_cli.service.list_training_jobs.ListTrainingJobs.list_training_jobs"
-    )
-    @mock.patch("logging.Logger.debug")
-    def test_list_job_happy_case_default_output_mode_json(
-        self,
-        mock_debug,
-        mock_list_training_job_service_and_list_jobs: mock.Mock,
-        mock_list_training_job_service: mock.Mock,
-    ):
-        mock_list_training_job_service.return_value = self.mock_list_jobs
-        mock_list_training_job_service_and_list_jobs.return_value = json.dumps({"jobs": [{
-            "Name": "test-job-name",
-            "Namespace": "test_namespace",
-            "CreationTime": "2025-01-01T01:01:01Z",
-            "State": "Succeeded"
-        }]})
-        result = self.runner.invoke(list_jobs, ["--output", "json"])
-        self.assertEqual(result.exit_code, 0)
-        expected_output = '{\n    "jobs": [\n        {\n            "Name": "test-job-name",\n            "Namespace": "test_namespace",\n            "CreationTime": "2025-01-01T01:01:01Z",\n            "State": "Succeeded"\n        }\n    ]\n}\n'
-        self.assertEqual(expected_output, result.output)
-        mock_debug.assert_called()
-
-    @mock.patch("hyperpod_cli.service.list_training_jobs.ListTrainingJobs")
-    @mock.patch(
-        "hyperpod_cli.service.list_training_jobs.ListTrainingJobs.list_training_jobs"
-    )
-    @mock.patch("logging.Logger.debug")
-    def test_list_job_happy_case_default_output_mode_table(
-        self,
-        mock_debug,
-        mock_list_training_job_service_and_list_jobs: mock.Mock,
-        mock_list_training_job_service: mock.Mock,
-    ):
-        mock_list_training_job_service.return_value = self.mock_list_jobs
-        mock_list_training_job_service_and_list_jobs.return_value = json.dumps({"jobs": [{
-            "Name": "test-job-name",
-            "Namespace": "test_namespace",
-            "CreationTime": "2025-01-01T01:01:01Z",
-            "State": "Succeeded"
-        }]})
-        result = self.runner.invoke(list_jobs, ["--output", "table"])
-        self.assertEqual(result.exit_code, 0)
-        expected_output = ' Name          | Namespace      | CreationTime         | State\n---------------+----------------+----------------------+-----------\n test-job-name | test_namespace | 2025-01-01T01:01:01Z | Succeeded\n'
-        self.assertEqual(expected_output, result.output)
-        mock_debug.assert_called()
-
     @mock.patch("hyperpod_cli.service.list_training_jobs.ListTrainingJobs")
     @mock.patch(
         "hyperpod_cli.service.list_training_jobs.ListTrainingJobs.list_training_jobs"
@@ -241,7 +192,7 @@ def test_list_job_happy_case_with_namespace(
         mock_list_training_job_service: mock.Mock,
     ):
         mock_list_training_job_service.return_value = self.mock_list_jobs
-        mock_list_training_job_service_and_list_jobs.return_value = json.dumps({"jobs": []})
+        mock_list_training_job_service_and_list_jobs.return_value = {"jobs": []}
         result = self.runner.invoke(list_jobs, ["--namespace", "kubeflow"])
         self.assertEqual(result.exit_code, 0)
         self.assertIn("jobs", result.output)
@@ -256,7 +207,7 @@ def test_list_job_happy_case_with_all_namespace(
         mock_list_training_job_service: mock.Mock,
     ):
         mock_list_training_job_service.return_value = self.mock_list_jobs
-        mock_list_training_job_service_and_list_jobs.return_value = json.dumps({'jobs': []})
+        mock_list_training_job_service_and_list_jobs.return_value = {"jobs": []}
         result = self.runner.invoke(list_jobs, ["-A"])
         self.assertEqual(result.exit_code, 0)
         self.assertIn("jobs", result.output)
@@ -271,7 +222,7 @@ def test_list_job_happy_case_with_all_namespace_and_selector(
         mock_list_training_job_service: mock.Mock,
     ):
         mock_list_training_job_service.return_value = self.mock_list_jobs
-        mock_list_training_job_service_and_list_jobs.return_value = json.dumps({"jobs": []})
+        mock_list_training_job_service_and_list_jobs.return_value = {"jobs": []}
         result = self.runner.invoke(list_jobs, ["-A", "-l", "test=test"])
         self.assertEqual(result.exit_code, 0)
         self.assertIn("jobs", result.output)