diff --git a/src/codeflare_sdk/ray/__init__.py b/src/codeflare_sdk/ray/__init__.py index 7c8e84da..b2278a05 100644 --- a/src/codeflare_sdk/ray/__init__.py +++ b/src/codeflare_sdk/ray/__init__.py @@ -6,6 +6,9 @@ from .rayjobs import ( RayJob, + RayJobDeploymentStatus, + CodeflareRayJobStatus, + RayJobInfo, ) from .cluster import ( diff --git a/src/codeflare_sdk/ray/rayjobs/__init__.py b/src/codeflare_sdk/ray/rayjobs/__init__.py index d9cbae34..47b573af 100644 --- a/src/codeflare_sdk/ray/rayjobs/__init__.py +++ b/src/codeflare_sdk/ray/rayjobs/__init__.py @@ -1 +1,2 @@ from .rayjob import RayJob +from .status import RayJobDeploymentStatus, CodeflareRayJobStatus, RayJobInfo diff --git a/src/codeflare_sdk/ray/rayjobs/pretty_print.py b/src/codeflare_sdk/ray/rayjobs/pretty_print.py new file mode 100644 index 00000000..9bc89b88 --- /dev/null +++ b/src/codeflare_sdk/ray/rayjobs/pretty_print.py @@ -0,0 +1,116 @@ +# Copyright 2025 IBM, Red Hat +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is 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. + +""" +This sub-module exists primarily to be used internally by the RayJob object +(in the rayjob sub-module) for pretty-printing job status and details. +""" + +from rich.console import Console +from rich.table import Table +from rich.panel import Panel +from typing import Tuple, Optional + +from .status import RayJobDeploymentStatus, RayJobInfo + + +def print_job_status(job_info: RayJobInfo): + """ + Pretty print the job status in a format similar to cluster status. + """ + status_display, header_color = _get_status_display(job_info.status) + + # Create main info table + table = _create_info_table(header_color, job_info.name, status_display) + table.add_row(f"[bold]Job ID:[/bold] {job_info.job_id}") + table.add_row(f"[bold]Status:[/bold] {job_info.status.value}") + table.add_row(f"[bold]RayCluster:[/bold] {job_info.cluster_name}") + table.add_row(f"[bold]Namespace:[/bold] {job_info.namespace}") + + # Add timing information if available + if job_info.start_time: + table.add_row(f"[bold]Started:[/bold] {job_info.start_time}") + + # Add attempt counts if there are failures + if job_info.failed_attempts > 0: + table.add_row(f"[bold]Failed Attempts:[/bold] {job_info.failed_attempts}") + + _print_table_in_panel(table) + + +def print_no_job_found(job_name: str, namespace: str): + """ + Print a message when no job is found. + """ + # Create table with error message + table = _create_info_table( + "[white on red][bold]Name", job_name, "[bold red]No RayJob found" + ) + table.add_row() + table.add_row("Please run rayjob.submit() to submit a job.") + table.add_row() + table.add_row(f"[bold]Namespace:[/bold] {namespace}") + + _print_table_in_panel(table) + + +def _get_status_display(status: RayJobDeploymentStatus) -> Tuple[str, str]: + """ + Get the display string and header color for a given status. + + Returns: + Tuple of (status_display, header_color) + """ + status_mapping = { + RayJobDeploymentStatus.COMPLETE: ( + "Complete :white_heavy_check_mark:", + "[white on green][bold]Name", + ), + RayJobDeploymentStatus.RUNNING: ("Running :gear:", "[white on blue][bold]Name"), + RayJobDeploymentStatus.FAILED: ("Failed :x:", "[white on red][bold]Name"), + RayJobDeploymentStatus.SUSPENDED: ( + "Suspended :pause_button:", + "[white on yellow][bold]Name", + ), + } + + return status_mapping.get( + status, ("Unknown :question:", "[white on red][bold]Name") + ) + + +def _create_info_table(header_color: str, name: str, status_display: str) -> Table: + """ + Create a standardized info table with header and status. + + Returns: + Table with header row, name/status row, and empty separator row + """ + table = Table(box=None, show_header=False) + table.add_row(header_color) + table.add_row("[bold underline]" + name, status_display) + table.add_row() # Empty separator row + return table + + +def _print_table_in_panel(table: Table): + """ + Print a table wrapped in a consistent panel format. + """ + console = Console() + main_table = Table( + box=None, title="[bold] :package: CodeFlare RayJob Status :package:" + ) + main_table.add_row(Panel.fit(table)) + console.print(main_table) diff --git a/src/codeflare_sdk/ray/rayjobs/rayjob.py b/src/codeflare_sdk/ray/rayjobs/rayjob.py index e7a9a588..ac2210a2 100644 --- a/src/codeflare_sdk/ray/rayjobs/rayjob.py +++ b/src/codeflare_sdk/ray/rayjobs/rayjob.py @@ -1,11 +1,32 @@ +# Copyright 2025 IBM, Red Hat +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is 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. + """ RayJob client for submitting and managing Ray jobs using the odh-kuberay-client. """ import logging -from typing import Dict, Any, Optional +from typing import Dict, Any, Optional, Tuple from odh_kuberay_client.kuberay_job_api import RayjobApi +from .status import ( + RayJobDeploymentStatus, + CodeflareRayJobStatus, + RayJobInfo, +) +from . import pretty_print + # Set up logging logger = logging.getLogger(__name__) @@ -15,7 +36,7 @@ class RayJob: A client for managing Ray jobs using the KubeRay operator. This class provides a simplified interface for submitting and managing - Ray jobs in a Kubernetes cluster with the KubeRay operator installed. + RayJob CRs (using the KubeRay RayJob python client). """ def __init__( @@ -109,3 +130,73 @@ def _build_rayjob_cr( rayjob_cr["spec"]["runtimeEnvYAML"] = str(runtime_env) return rayjob_cr + + def status( + self, print_to_console: bool = True + ) -> Tuple[CodeflareRayJobStatus, bool]: + """ + Get the status of the Ray job. + + Args: + print_to_console (bool): Whether to print formatted status to console (default: True) + + Returns: + Tuple of (CodeflareRayJobStatus, ready: bool) where ready indicates job completion + """ + status_data = self._api.get_job_status( + name=self.name, k8s_namespace=self.namespace + ) + + if not status_data: + if print_to_console: + pretty_print.print_no_job_found(self.name, self.namespace) + return CodeflareRayJobStatus.UNKNOWN, False + + # Map deployment status to our enums + deployment_status_str = status_data.get("jobDeploymentStatus", "Unknown") + + try: + deployment_status = RayJobDeploymentStatus(deployment_status_str) + except ValueError: + deployment_status = RayJobDeploymentStatus.UNKNOWN + + # Create RayJobInfo dataclass + job_info = RayJobInfo( + name=self.name, + job_id=status_data.get("jobId", ""), + status=deployment_status, + namespace=self.namespace, + cluster_name=self.cluster_name, + start_time=status_data.get("startTime"), + end_time=status_data.get("endTime"), + failed_attempts=status_data.get("failed", 0), + succeeded_attempts=status_data.get("succeeded", 0), + ) + + # Map to CodeFlare status and determine readiness + codeflare_status, ready = self._map_to_codeflare_status(deployment_status) + + if print_to_console: + pretty_print.print_job_status(job_info) + + return codeflare_status, ready + + def _map_to_codeflare_status( + self, deployment_status: RayJobDeploymentStatus + ) -> Tuple[CodeflareRayJobStatus, bool]: + """ + Map deployment status to CodeFlare status and determine readiness. + + Returns: + Tuple of (CodeflareRayJobStatus, ready: bool) + """ + status_mapping = { + RayJobDeploymentStatus.COMPLETE: (CodeflareRayJobStatus.COMPLETE, True), + RayJobDeploymentStatus.RUNNING: (CodeflareRayJobStatus.RUNNING, False), + RayJobDeploymentStatus.FAILED: (CodeflareRayJobStatus.FAILED, False), + RayJobDeploymentStatus.SUSPENDED: (CodeflareRayJobStatus.SUSPENDED, False), + } + + return status_mapping.get( + deployment_status, (CodeflareRayJobStatus.UNKNOWN, False) + ) diff --git a/src/codeflare_sdk/ray/rayjobs/status.py b/src/codeflare_sdk/ray/rayjobs/status.py new file mode 100644 index 00000000..027ed09c --- /dev/null +++ b/src/codeflare_sdk/ray/rayjobs/status.py @@ -0,0 +1,64 @@ +# Copyright 2025 IBM, Red Hat +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is 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. + +""" +The status sub-module defines Enums containing information for Ray job +deployment states and CodeFlare job states, as well as +dataclasses to store information for Ray jobs. +""" + +from dataclasses import dataclass +from enum import Enum +from typing import Optional + + +class RayJobDeploymentStatus(Enum): + """ + Defines the possible deployment states of a Ray job (from the KubeRay RayJob API). + """ + + COMPLETE = "Complete" + RUNNING = "Running" + FAILED = "Failed" + SUSPENDED = "Suspended" + UNKNOWN = "Unknown" + + +class CodeflareRayJobStatus(Enum): + """ + Defines the possible reportable states of a CodeFlare Ray job. + """ + + COMPLETE = 1 + RUNNING = 2 + FAILED = 3 + SUSPENDED = 4 + UNKNOWN = 5 + + +@dataclass +class RayJobInfo: + """ + For storing information about a Ray job. + """ + + name: str + job_id: str + status: RayJobDeploymentStatus + namespace: str + cluster_name: str + start_time: Optional[str] = None + end_time: Optional[str] = None + failed_attempts: int = 0 + succeeded_attempts: int = 0 diff --git a/src/codeflare_sdk/ray/rayjobs/test_pretty_print.py b/src/codeflare_sdk/ray/rayjobs/test_pretty_print.py new file mode 100644 index 00000000..dbfd7caf --- /dev/null +++ b/src/codeflare_sdk/ray/rayjobs/test_pretty_print.py @@ -0,0 +1,262 @@ +# Copyright 2025 IBM, Red Hat +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is 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. + +from codeflare_sdk.ray.rayjobs.pretty_print import ( + _get_status_display, + print_job_status, + print_no_job_found, +) +from codeflare_sdk.ray.rayjobs.status import RayJobDeploymentStatus, RayJobInfo +from unittest.mock import MagicMock, call + + +def test_get_status_display(): + """ + Test the _get_status_display function. + """ + # Test Complete status + display, color = _get_status_display(RayJobDeploymentStatus.COMPLETE) + assert display == "Complete :white_heavy_check_mark:" + assert color == "[white on green][bold]Name" + + # Test Running status + display, color = _get_status_display(RayJobDeploymentStatus.RUNNING) + assert display == "Running :gear:" + assert color == "[white on blue][bold]Name" + + # Test Failed status + display, color = _get_status_display(RayJobDeploymentStatus.FAILED) + assert display == "Failed :x:" + assert color == "[white on red][bold]Name" + + # Test Suspended status + display, color = _get_status_display(RayJobDeploymentStatus.SUSPENDED) + assert display == "Suspended :pause_button:" + assert color == "[white on yellow][bold]Name" + + # Test Unknown status + display, color = _get_status_display(RayJobDeploymentStatus.UNKNOWN) + assert display == "Unknown :question:" + assert color == "[white on red][bold]Name" + + +def test_print_job_status_running_format(mocker): + """ + Test the print_job_status function format for a running job. + """ + # Mock Rich components to verify format + mock_console = MagicMock() + mock_inner_table = MagicMock() + mock_main_table = MagicMock() + mock_panel = MagicMock() + + # Mock Table to return different instances for inner and main tables + table_instances = [mock_inner_table, mock_main_table] + mock_table_class = MagicMock(side_effect=table_instances) + + mocker.patch( + "codeflare_sdk.ray.rayjobs.pretty_print.Console", return_value=mock_console + ) + mocker.patch("codeflare_sdk.ray.rayjobs.pretty_print.Table", mock_table_class) + mocker.patch("codeflare_sdk.ray.rayjobs.pretty_print.Panel", mock_panel) + + # Create test job info for running job + job_info = RayJobInfo( + name="test-job", + job_id="test-job-abc123", + status=RayJobDeploymentStatus.RUNNING, + namespace="test-ns", + cluster_name="test-cluster", + start_time="2025-07-28T11:37:07Z", + failed_attempts=0, + succeeded_attempts=0, + ) + + # Call the function + print_job_status(job_info) + + # Verify both Table calls + expected_table_calls = [ + call(box=None, show_header=False), # Inner content table + call( + box=None, title="[bold] :package: CodeFlare RayJob Status :package:" + ), # Main wrapper table + ] + mock_table_class.assert_has_calls(expected_table_calls) + + # Verify inner table rows are added in correct order and format (versus our hard-coded version of this for cluster) + expected_calls = [ + call("[white on blue][bold]Name"), # Header with blue color for running + call( + "[bold underline]test-job", "Running :gear:" + ), # Name and status with gear emoji + call(), # Empty separator row + call("[bold]Job ID:[/bold] test-job-abc123"), + call("[bold]Status:[/bold] Running"), + call("[bold]RayCluster:[/bold] test-cluster"), + call("[bold]Namespace:[/bold] test-ns"), + call("[bold]Started:[/bold] 2025-07-28T11:37:07Z"), + ] + mock_inner_table.add_row.assert_has_calls(expected_calls) + + # Verify Panel is created with inner table + mock_panel.fit.assert_called_once_with(mock_inner_table) + + # Verify main table gets the panel + mock_main_table.add_row.assert_called_once_with(mock_panel.fit.return_value) + + # Verify console prints the main table + mock_console.print.assert_called_once_with(mock_main_table) + + +def test_print_job_status_complete_format(mocker): + """ + Test the print_job_status function format for a completed job. + """ + # Mock Rich components + mock_console = MagicMock() + mock_inner_table = MagicMock() + mock_main_table = MagicMock() + mock_panel = MagicMock() + + # Mock Table to return different instances + table_instances = [mock_inner_table, mock_main_table] + mock_table_class = MagicMock(side_effect=table_instances) + + mocker.patch( + "codeflare_sdk.ray.rayjobs.pretty_print.Console", return_value=mock_console + ) + mocker.patch("codeflare_sdk.ray.rayjobs.pretty_print.Table", mock_table_class) + mocker.patch("codeflare_sdk.ray.rayjobs.pretty_print.Panel", mock_panel) + + # Create test job info for completed job + job_info = RayJobInfo( + name="completed-job", + job_id="completed-job-xyz789", + status=RayJobDeploymentStatus.COMPLETE, + namespace="prod-ns", + cluster_name="prod-cluster", + start_time="2025-07-28T11:37:07Z", + failed_attempts=0, + succeeded_attempts=1, + ) + + # Call the function + print_job_status(job_info) + + # Verify correct header color for completed job (green) (versus our hard-coded version of this for cluster) + expected_calls = [ + call("[white on green][bold]Name"), # Green header for complete + call( + "[bold underline]completed-job", "Complete :white_heavy_check_mark:" + ), # Checkmark emoji + call(), # Empty separator + call("[bold]Job ID:[/bold] completed-job-xyz789"), + call("[bold]Status:[/bold] Complete"), + call("[bold]RayCluster:[/bold] prod-cluster"), + call("[bold]Namespace:[/bold] prod-ns"), + call("[bold]Started:[/bold] 2025-07-28T11:37:07Z"), + ] + mock_inner_table.add_row.assert_has_calls(expected_calls) + + +def test_print_job_status_failed_with_attempts_format(mocker): + """ + Test the print_job_status function format for a failed job with attempts. + """ + # Mock Rich components + mock_console = MagicMock() + mock_inner_table = MagicMock() + mock_main_table = MagicMock() + mock_panel = MagicMock() + + # Mock Table to return different instances + table_instances = [mock_inner_table, mock_main_table] + mock_table_class = MagicMock(side_effect=table_instances) + + mocker.patch( + "codeflare_sdk.ray.rayjobs.pretty_print.Console", return_value=mock_console + ) + mocker.patch("codeflare_sdk.ray.rayjobs.pretty_print.Table", mock_table_class) + mocker.patch("codeflare_sdk.ray.rayjobs.pretty_print.Panel", mock_panel) + + # Create test job info with failures + job_info = RayJobInfo( + name="failing-job", + job_id="failing-job-fail123", + status=RayJobDeploymentStatus.FAILED, + namespace="test-ns", + cluster_name="test-cluster", + start_time="2025-07-28T11:37:07Z", + failed_attempts=3, # Has failures + succeeded_attempts=0, + ) + + # Call the function + print_job_status(job_info) + + # Verify correct formatting including failure attempts (versus our hard-coded version of this for cluster) + expected_calls = [ + call("[white on red][bold]Name"), # Red header for failed + call("[bold underline]failing-job", "Failed :x:"), # X emoji for failed + call(), # Empty separator + call("[bold]Job ID:[/bold] failing-job-fail123"), + call("[bold]Status:[/bold] Failed"), + call("[bold]RayCluster:[/bold] test-cluster"), + call("[bold]Namespace:[/bold] test-ns"), + call("[bold]Started:[/bold] 2025-07-28T11:37:07Z"), + call("[bold]Failed Attempts:[/bold] 3"), # Failed attempts should be shown + ] + mock_inner_table.add_row.assert_has_calls(expected_calls) + + +def test_print_no_job_found_format(mocker): + """ + Test the print_no_job_found function format. + """ + # Mock Rich components + mock_console = MagicMock() + mock_inner_table = MagicMock() + mock_main_table = MagicMock() + mock_panel = MagicMock() + + # Mock Table to return different instances + table_instances = [mock_inner_table, mock_main_table] + mock_table_class = MagicMock(side_effect=table_instances) + + mocker.patch( + "codeflare_sdk.ray.rayjobs.pretty_print.Console", return_value=mock_console + ) + mocker.patch("codeflare_sdk.ray.rayjobs.pretty_print.Table", mock_table_class) + mocker.patch("codeflare_sdk.ray.rayjobs.pretty_print.Panel", mock_panel) + + # Call the function + print_no_job_found("missing-job", "test-namespace") + + # Verify error message format (versus our hard-coded version of this for cluster) + expected_calls = [ + call("[white on red][bold]Name"), # Red header for error + call( + "[bold underline]missing-job", "[bold red]No RayJob found" + ), # Error message in red + call(), # Empty separator + call(), # Another empty row + call("Please run rayjob.submit() to submit a job."), # Helpful hint + call(), # Empty separator + call("[bold]Namespace:[/bold] test-namespace"), + ] + mock_inner_table.add_row.assert_has_calls(expected_calls) + + # Verify Panel is used + mock_panel.fit.assert_called_once_with(mock_inner_table) diff --git a/src/codeflare_sdk/ray/rayjobs/test_rayjob.py b/src/codeflare_sdk/ray/rayjobs/test_rayjob.py index 1136e6e5..5429f303 100644 --- a/src/codeflare_sdk/ray/rayjobs/test_rayjob.py +++ b/src/codeflare_sdk/ray/rayjobs/test_rayjob.py @@ -1,4 +1,4 @@ -# Copyright 2024 IBM, Red Hat +# Copyright 2025 IBM, Red Hat # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/src/codeflare_sdk/ray/rayjobs/test_status.py b/src/codeflare_sdk/ray/rayjobs/test_status.py new file mode 100644 index 00000000..6d2ce946 --- /dev/null +++ b/src/codeflare_sdk/ray/rayjobs/test_status.py @@ -0,0 +1,290 @@ +# Copyright 2025 IBM, Red Hat +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is 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. + +from codeflare_sdk.ray.rayjobs.rayjob import RayJob +from codeflare_sdk.ray.rayjobs.status import ( + CodeflareRayJobStatus, + RayJobDeploymentStatus, + RayJobInfo, +) + + +def test_rayjob_status(mocker): + """ + Test the RayJob status method with different deployment statuses. + """ + # Mock the RayjobApi to avoid actual Kubernetes calls + mock_api_class = mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi") + mock_api_instance = mock_api_class.return_value + + # Create a RayJob instance + rayjob = RayJob( + job_name="test-job", + cluster_name="test-cluster", + namespace="test-ns", + entrypoint="python test.py", + ) + + # Test case 1: No job found + mock_api_instance.get_job_status.return_value = None + status, ready = rayjob.status(print_to_console=False) + assert status == CodeflareRayJobStatus.UNKNOWN + assert ready == False + + # Test case 2: Running job + mock_api_instance.get_job_status.return_value = { + "jobId": "test-job-abc123", + "jobDeploymentStatus": "Running", + "startTime": "2025-07-28T11:37:07Z", + "failed": 0, + "succeeded": 0, + "rayClusterName": "test-cluster", + } + status, ready = rayjob.status(print_to_console=False) + assert status == CodeflareRayJobStatus.RUNNING + assert ready == False + + # Test case 3: Complete job + mock_api_instance.get_job_status.return_value = { + "jobId": "test-job-abc123", + "jobDeploymentStatus": "Complete", + "startTime": "2025-07-28T11:37:07Z", + "endTime": "2025-07-28T11:42:30Z", + "failed": 0, + "succeeded": 1, + "rayClusterName": "test-cluster", + } + status, ready = rayjob.status(print_to_console=False) + assert status == CodeflareRayJobStatus.COMPLETE + assert ready == True + + # Test case 4: Failed job + mock_api_instance.get_job_status.return_value = { + "jobId": "test-job-abc123", + "jobDeploymentStatus": "Failed", + "startTime": "2025-07-28T11:37:07Z", + "endTime": "2025-07-28T11:42:30Z", + "failed": 1, + "succeeded": 0, + "rayClusterName": "test-cluster", + } + status, ready = rayjob.status(print_to_console=False) + assert status == CodeflareRayJobStatus.FAILED + assert ready == False + + # Test case 5: Suspended job + mock_api_instance.get_job_status.return_value = { + "jobId": "test-job-abc123", + "jobDeploymentStatus": "Suspended", + "startTime": "2025-07-28T11:37:07Z", + "failed": 0, + "succeeded": 0, + "rayClusterName": "test-cluster", + } + status, ready = rayjob.status(print_to_console=False) + assert status == CodeflareRayJobStatus.SUSPENDED + assert ready == False + + +def test_rayjob_status_unknown_deployment_status(mocker): + """ + Test handling of unknown deployment status from the API. + """ + mock_api_class = mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi") + mock_api_instance = mock_api_class.return_value + + rayjob = RayJob( + job_name="test-job", + cluster_name="test-cluster", + namespace="test-ns", + entrypoint="python test.py", + ) + + # Test with unrecognized deployment status + mock_api_instance.get_job_status.return_value = { + "jobId": "test-job-abc123", + "jobDeploymentStatus": "SomeNewStatus", # Unknown status + "startTime": "2025-07-28T11:37:07Z", + "failed": 0, + "succeeded": 0, + } + + status, ready = rayjob.status(print_to_console=False) + assert status == CodeflareRayJobStatus.UNKNOWN + assert ready == False + + +def test_rayjob_status_missing_fields(mocker): + """ + Test handling of API response with missing fields. + """ + mock_api_class = mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi") + mock_api_instance = mock_api_class.return_value + + rayjob = RayJob( + job_name="test-job", + cluster_name="test-cluster", + namespace="test-ns", + entrypoint="python test.py", + ) + + # Test with minimal API response (missing some fields) + mock_api_instance.get_job_status.return_value = { + # Missing jobId, failed, succeeded, etc. + "jobDeploymentStatus": "Running", + } + + status, ready = rayjob.status(print_to_console=False) + assert status == CodeflareRayJobStatus.RUNNING + assert ready == False + + +def test_map_to_codeflare_status(mocker): + """ + Test the _map_to_codeflare_status helper method directly. + """ + # Mock the RayjobApi constructor to avoid authentication issues + mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi") + + rayjob = RayJob( + job_name="test-job", + cluster_name="test-cluster", + namespace="test-ns", + entrypoint="python test.py", + ) + + # Test all deployment status mappings + status, ready = rayjob._map_to_codeflare_status(RayJobDeploymentStatus.COMPLETE) + assert status == CodeflareRayJobStatus.COMPLETE + assert ready == True + + status, ready = rayjob._map_to_codeflare_status(RayJobDeploymentStatus.RUNNING) + assert status == CodeflareRayJobStatus.RUNNING + assert ready == False + + status, ready = rayjob._map_to_codeflare_status(RayJobDeploymentStatus.FAILED) + assert status == CodeflareRayJobStatus.FAILED + assert ready == False + + status, ready = rayjob._map_to_codeflare_status(RayJobDeploymentStatus.SUSPENDED) + assert status == CodeflareRayJobStatus.SUSPENDED + assert ready == False + + status, ready = rayjob._map_to_codeflare_status(RayJobDeploymentStatus.UNKNOWN) + assert status == CodeflareRayJobStatus.UNKNOWN + assert ready == False + + +def test_rayjob_info_dataclass(): + """ + Test the RayJobInfo dataclass creation and field access. + """ + job_info = RayJobInfo( + name="test-job", + job_id="test-job-abc123", + status=RayJobDeploymentStatus.RUNNING, + namespace="test-ns", + cluster_name="test-cluster", + start_time="2025-07-28T11:37:07Z", + failed_attempts=0, + succeeded_attempts=0, + ) + + # Test all fields are accessible + assert job_info.name == "test-job" + assert job_info.job_id == "test-job-abc123" + assert job_info.status == RayJobDeploymentStatus.RUNNING + assert job_info.namespace == "test-ns" + assert job_info.cluster_name == "test-cluster" + assert job_info.start_time == "2025-07-28T11:37:07Z" + assert job_info.end_time is None # Default value + assert job_info.failed_attempts == 0 + assert job_info.succeeded_attempts == 0 + + +def test_rayjob_status_print_no_job_found(mocker): + """ + Test that pretty_print.print_no_job_found is called when no job is found and print_to_console=True. + """ + # Mock the RayjobApi and pretty_print + mock_api_class = mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi") + mock_api_instance = mock_api_class.return_value + mock_print_no_job_found = mocker.patch( + "codeflare_sdk.ray.rayjobs.pretty_print.print_no_job_found" + ) + + # Create a RayJob instance + rayjob = RayJob( + job_name="test-job", + cluster_name="test-cluster", + namespace="test-ns", + entrypoint="python test.py", + ) + + # No job found scenario + mock_api_instance.get_job_status.return_value = None + + # Call status with print_to_console=True + status, ready = rayjob.status(print_to_console=True) + + # Verify the pretty print function was called + mock_print_no_job_found.assert_called_once_with("test-job", "test-ns") + assert status == CodeflareRayJobStatus.UNKNOWN + assert ready == False + + +def test_rayjob_status_print_job_found(mocker): + """ + Test that pretty_print.print_job_status is called when job is found and print_to_console=True. + """ + # Mock the RayjobApi and pretty_print + mock_api_class = mocker.patch("codeflare_sdk.ray.rayjobs.rayjob.RayjobApi") + mock_api_instance = mock_api_class.return_value + mock_print_job_status = mocker.patch( + "codeflare_sdk.ray.rayjobs.pretty_print.print_job_status" + ) + + # Create a RayJob instance + rayjob = RayJob( + job_name="test-job", + cluster_name="test-cluster", + namespace="test-ns", + entrypoint="python test.py", + ) + + # Job found scenario + mock_api_instance.get_job_status.return_value = { + "jobId": "test-job-abc123", + "jobDeploymentStatus": "Running", + "startTime": "2025-07-28T11:37:07Z", + "failed": 0, + "succeeded": 0, + "rayClusterName": "test-cluster", + } + + # Call status with print_to_console=True + status, ready = rayjob.status(print_to_console=True) + + # Verify the pretty print function was called + mock_print_job_status.assert_called_once() + # Verify the RayJobInfo object passed to print_job_status + call_args = mock_print_job_status.call_args[0][0] # First positional argument + assert call_args.name == "test-job" + assert call_args.job_id == "test-job-abc123" + assert call_args.status == RayJobDeploymentStatus.RUNNING + assert call_args.namespace == "test-ns" + assert call_args.cluster_name == "test-cluster" + + assert status == CodeflareRayJobStatus.RUNNING + assert ready == False