bl only print logs for kueue specific pods + unit tests

Author: memalhot

batchtools/bl.py

@@ -8,6 +8,7 @@
 from .basecommand import Command
 from .basecommand import SubParserFactory
 from .helpers import pretty_print
+from .helpers import is_kueue_managed_pod
 
 
 class LogsCommandArgs(argparse.Namespace):
@@ -63,10 +64,15 @@ def run(args: argparse.Namespace):
                     print(pretty_print(pod_dict[name]))
 
             else:
-                # case where user provides no args, print logs for all pods
+                printed_any = False
                 for name, pod in pod_dict.items():
-                    print(f"\nLogs for {name}:\n{'-' * 40}")
-                    print(pretty_print(pod))
+                    if is_kueue_managed_pod(pod):
+                        printed_any = True
+                        print(f"\nLogs for {name}:\n{'-' * 40}")
+                        print(pretty_print(pod))
+
+                if not printed_any:
+                    print("No Kueue-managed pods found")
 
         except oc.OpenShiftPythonException as e:
             sys.exit(f"Error occurred while retrieving logs: {e}")

batchtools/helpers.py

@@ -37,3 +37,18 @@ def is_kueue_managed_job(job_obj) -> bool:
             return True
     except Exception:
         return False
+
+
+def is_kueue_managed_pod(pod) -> bool:
+    try:
+        owners = getattr(pod.model.metadata, "ownerReferences", []) or []
+        job_owner = next((o for o in owners if o.kind == "Job"), None)
+        if not job_owner:
+            return False
+
+        job_name = job_owner.name
+        job_obj = oc.selector(f"job/{job_name}").object()
+        return is_kueue_managed_job(job_obj)
+
+    except Exception:
+        return False

tests/test_bl.py

@@ -1,6 +1,7 @@
 import pytest
 from unittest import mock
 from contextlib import contextmanager
+import openshift_client as oc
 
 import argparse
 from typing import Any
@@ -41,7 +42,10 @@ def args() -> argparse.Namespace:
 
 @contextmanager
 def patch_pods_selector(pods: list[Any]):
-    with mock.patch("openshift_client.selector") as mock_selector:
+    with (
+        mock.patch("openshift_client.selector") as mock_selector,
+        mock.patch("batchtools.bl.is_kueue_managed_pod", return_value=True),
+    ):
         mock_result = mock.Mock(name="result")
         mock_result.objects.return_value = pods
         mock_selector.return_value = mock_result
@@ -75,3 +79,45 @@ def test_get_logs_selected(args: argparse.Namespace, pods: list[Any], capsys):
         captured = capsys.readouterr()
         assert "Logs for pod1" in captured.out
         assert "Logs for pod2" not in captured.out
+
+
+def test_no_kueue_managed_pods(args: argparse.Namespace, pods: list[Any], capsys):
+    with (
+        patch_pods_selector(pods),
+        mock.patch("batchtools.bl.is_kueue_managed_pod", return_value=False),
+    ):
+        args.pod_names = []
+        LogsCommand.run(args)
+        captured = capsys.readouterr()
+        assert "No Kueue-managed pods found" in captured.out
+
+
+def test_invalid_pod_name(args: argparse.Namespace, pods: list[Any], capsys):
+    with patch_pods_selector(pods):
+        args.pod_names = ["pod1", "does-not-exist"]
+        LogsCommand.run(args)
+        captured = capsys.readouterr()
+
+        # Valid pod still prints logs
+        assert "Logs for pod1" in captured.out
+        # Invalid pod hits the error path
+        assert (
+            "does-not-exist is not a valid pod. Logs cannot be retrieved."
+            in captured.out
+        )
+
+
+def test_selector_exception_exits_cleanly(args: argparse.Namespace):
+    # Make openshift_client.selector raise the same exception type used in bl.py
+    with mock.patch("openshift_client.selector") as mock_selector:
+        mock_selector.side_effect = oc.OpenShiftPythonException(
+            "test exception",
+            mock.Mock(),
+        )
+
+        with pytest.raises(SystemExit) as excinfo:
+            LogsCommand.run(args)
+
+        message = str(excinfo.value)
+        assert "Error occurred while retrieving logs:" in message
+        assert "test exception" in message