From 51c277d2ecef16869667e83c8c74249067775e33 Mon Sep 17 00:00:00 2001 From: Punit Naik Date: Wed, 3 Dec 2025 12:41:50 +0530 Subject: [PATCH 1/2] removed duplicate job status command --- chuck_data/commands/__init__.py | 2 +- chuck_data/commands/job_status.py | 6 +- chuck_data/commands/jobs.py | 68 +------------- tests/unit/commands/test_jobs.py | 146 +----------------------------- 4 files changed, 7 insertions(+), 215 deletions(-) diff --git a/chuck_data/commands/__init__.py b/chuck_data/commands/__init__.py index 52b3e0a..91e3341 100644 --- a/chuck_data/commands/__init__.py +++ b/chuck_data/commands/__init__.py @@ -85,7 +85,7 @@ setup_stitch_definition, add_stitch_report_definition, # Job commands - *jobs_definition, + jobs_definition, job_status_definition, list_jobs_definition, monitor_job_definition, diff --git a/chuck_data/commands/job_status.py b/chuck_data/commands/job_status.py index dfb2ff5..5484d20 100644 --- a/chuck_data/commands/job_status.py +++ b/chuck_data/commands/job_status.py @@ -536,11 +536,11 @@ def handle_command( "If no parameters provided, uses the last cached job ID.", handler=handle_command, parameters={ - "job-id": { + "job_id": { "type": "string", "description": "Chuck job ID (primary parameter).", }, - "run-id": { + "run_id": { "type": "string", "description": "Databricks run ID (legacy fallback).", }, @@ -554,7 +554,7 @@ def handle_command( needs_api_client=True, visible_to_user=True, visible_to_agent=True, - usage_hint="Usage: /job-status [--job_id ] [--live] OR /job-status --run_id OR /job-status (uses cached ID)", + usage_hint="Usage: /job-status [--job-id ] [--live] OR /job-status --run-id OR /job-status (uses cached ID)", condensed_action="Checking job status", ) diff --git a/chuck_data/commands/jobs.py b/chuck_data/commands/jobs.py index cbde272..a282303 100644 --- a/chuck_data/commands/jobs.py +++ b/chuck_data/commands/jobs.py @@ -64,57 +64,6 @@ def handle_launch_job(client: Optional[DatabricksAPIClient], **kwargs) -> Comman return CommandResult(False, error=e, message=str(e)) -def handle_job_status(client: Optional[DatabricksAPIClient], **kwargs) -> CommandResult: - """Query Databricks for the status of a one-time run. - - Args: - client: API client instance - **kwargs: run_id (str), tool_output_callback (callable, optional) - """ - run_id_str: str = kwargs.get("run_id") - tool_output_callback = kwargs.get("tool_output_callback") - - if not run_id_str: - return CommandResult(False, message="run_id parameter is required.") - if not client: - return CommandResult(False, message="Client required to get job status.") - try: - if tool_output_callback: - tool_output_callback( - "job_status_progress", - {"step": f"Attempting to get status for run ID {run_id_str}."}, - ) - - data = client.get_job_run_status(run_id_str) - status = data.get("status", data.get("state", {})) - life_cycle_state = status.get("life_cycle_state", status.get("state")) - result_state = status.get( - "result_state", status.get("termination_details", {}).get("code") - ) - state_message = status.get( - "state_message", status.get("termination_details", {}).get("message") - ) - run_page_url = data.get("run_page_url") - msg = f"Run {run_id_str}: Status: {life_cycle_state or 'N/A'}, Result: {result_state or 'N/A'}, Message: {state_message or ''}" - if run_page_url: - msg += f" URL: {run_page_url}" - - if tool_output_callback: - tool_output_callback( - "job_status_progress", - { - "step": f"Status for run ID {run_id_str} retrieved.", - "status_data": data, - }, - ) - return CommandResult(True, data=data, message=msg) - except Exception as e: - logging.error( - f"Failed to get job status for run '{run_id_str}': {e}", exc_info=True - ) - return CommandResult(False, error=e, message=str(e)) - - LAUNCH_JOB_DEFINITION = CommandDefinition( name="launch_job", description="Launch a Databricks job using a config file", @@ -141,19 +90,4 @@ def handle_job_status(client: Optional[DatabricksAPIClient], **kwargs) -> Comman tui_aliases=["launch-job"], ) -JOB_STATUS_DEFINITION = CommandDefinition( - name="job_status", - description="Get the status of a Databricks job run", - usage_hint="job_status --run_id=123456", - parameters={ - "run_id": {"type": "string", "description": "ID of the job run to check"}, - }, - required_params=["run_id"], - handler=handle_job_status, - needs_api_client=True, - visible_to_agent=True, - tui_aliases=["job-status"], -) - -# Combined definition for module -DEFINITION = [LAUNCH_JOB_DEFINITION, JOB_STATUS_DEFINITION] +DEFINITION = LAUNCH_JOB_DEFINITION diff --git a/tests/unit/commands/test_jobs.py b/tests/unit/commands/test_jobs.py index 9338f3a..77ea91d 100644 --- a/tests/unit/commands/test_jobs.py +++ b/tests/unit/commands/test_jobs.py @@ -1,11 +1,11 @@ """ -Tests for job-related command handlers (/launch_job, /job_status). +Tests for job-related command handlers (/launch_job). Behavioral tests focused on command execution patterns, aligned with CLAUDE.MD. """ from unittest.mock import patch -from chuck_data.commands.jobs import handle_launch_job, handle_job_status +from chuck_data.commands.jobs import handle_launch_job from chuck_data.commands.base import CommandResult from chuck_data.agent.tool_executor import execute_tool @@ -49,19 +49,6 @@ def test_direct_command_launch_job_failure_missing_init_script_path_parameter( ) -def test_direct_command_job_status_failure_missing_run_id_parameter( - databricks_client_stub, temp_config -): - """Test getting job status with missing run_id parameter.""" - with patch("chuck_data.config._config_manager", temp_config): - # run_id is intentionally omitted - result = handle_job_status(databricks_client_stub) - assert not result.success - assert ( - "run_id" in result.message.lower() or "parameter" in result.message.lower() - ) - - # --- Direct Command Execution Tests: handle_launch_job --- @@ -209,51 +196,6 @@ def capture_progress(tool_name, data): assert call_args["policy_id"] == "POLICY123ABC" -# --- Direct Command Execution Tests: handle_job_status --- - - -def test_handle_job_status_basic_success(databricks_client_stub, temp_config): - """Test getting job status with successful response.""" - with patch("chuck_data.config._config_manager", temp_config): - result = handle_job_status(databricks_client_stub, run_id="123456") - assert result.success - assert result.data["state"]["life_cycle_state"] == "RUNNING" - assert result.data["run_id"] == 123456 - - -def test_handle_job_status_http_error(databricks_client_stub, temp_config): - """Test getting job status with HTTP error response.""" - with patch("chuck_data.config._config_manager", temp_config): - - def get_status_failing(run_id): - raise Exception("Not Found") - - databricks_client_stub.get_job_run_status = get_status_failing - result = handle_job_status(databricks_client_stub, run_id="999999") - assert not result.success - assert "Not Found" in result.message - - -def test_handle_job_status_missing_token(temp_config): - """Test getting job status with missing API token (results in no client).""" - with patch("chuck_data.config._config_manager", temp_config): - result = handle_job_status( - None, run_id="123456" - ) # Simulates client not being initializable - assert not result.success - assert "Client required" in result.message - - -def test_handle_job_status_missing_url(temp_config): - """Test getting job status with missing workspace URL (results in no client).""" - with patch("chuck_data.config._config_manager", temp_config): - result = handle_job_status( - None, run_id="123456" - ) # Simulates client not being initializable - assert not result.success - assert "Client required" in result.message - - # --- Agent-Specific Behavioral Tests: handle_launch_job --- @@ -352,62 +294,6 @@ def failing_callback(tool_name, data): assert "Agent display system crashed" in result.message -# --- Agent-Specific Behavioral Tests: handle_job_status --- - - -def test_agent_job_status_success_shows_progress_steps( - databricks_client_stub, temp_config -): - """ - AGENT TEST: Getting job status successfully shows expected progress steps. - """ - with patch("chuck_data.config._config_manager", temp_config): - progress_steps = [] - - def capture_progress(tool_name, data): - assert tool_name == "job_status_progress" - progress_steps.append(data.get("step", str(data))) - - # Ensure the default stub behavior provides a valid status for this success test - # (The fixture default is RUNNING for run_id 123456) - - result = handle_job_status( - databricks_client_stub, - run_id="123456", - tool_output_callback=capture_progress, - ) - assert result.success is True - assert result.data["state"]["life_cycle_state"] == "RUNNING" - assert len(progress_steps) == 2, "Expected two progress steps." - assert progress_steps[0] == "Attempting to get status for run ID 123456." - assert progress_steps[1] == "Status for run ID 123456 retrieved." - - -def test_agent_job_status_callback_errors_bubble_up( - databricks_client_stub, temp_config -): - """ - AGENT TEST: Errors from tool_output_callback should result in a failed CommandResult for job status. - """ - with patch("chuck_data.config._config_manager", temp_config): - - def failing_callback(tool_name, data): - assert tool_name == "job_status_progress" - raise Exception("Agent display system crashed during status") - - # No need to change get_job_run_status for this test, as the first callback should fail. - - result = handle_job_status( - databricks_client_stub, - run_id="123456", - tool_output_callback=failing_callback, - ) - assert ( - not result.success - ), "Handler did not set success=False when callback failed during status." - assert "Agent display system crashed during status" in result.message - - # --- Agent Tool Executor Integration Tests --- @@ -468,31 +354,3 @@ def mock_submit_job_run( # Verify policy_id was passed through the agent tool executor assert len(captured_policy_id) == 1 assert captured_policy_id[0] == "AGENT_POLICY_ID_123" - - -def test_agent_tool_executor_job_status_integration( - databricks_client_stub, temp_config -): - """AGENT TEST: End-to-end integration for getting job status via execute_tool. - Assumes 'job_status' is the correct registered tool name. - """ - with patch("chuck_data.config._config_manager", temp_config): - databricks_client_stub.get_job_run_status = lambda run_id: { - "run_id": int(run_id), - "state": { - "life_cycle_state": "TERMINATED", - "result_state": "SUCCESS", - "state_message": "Job finished.", - }, - "job_id": 789, - "task_run_stats": {}, - } - tool_args = {"run_id": "777888"} - agent_result = execute_tool( - api_client=databricks_client_stub, - tool_name="job_status", - tool_args=tool_args, - ) - assert agent_result is not None - assert agent_result.get("run_id") == 777888 - assert agent_result.get("life_cycle_state") == "TERMINATED" From 04c7c402f42bbcbfc0b93960774f88fc00458f4e Mon Sep 17 00:00:00 2001 From: Punit Naik Date: Wed, 3 Dec 2025 15:16:00 +0530 Subject: [PATCH 2/2] mocked job cache fn for unit tests --- tests/conftest.py | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index de6abab..fafdf1f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -3,7 +3,7 @@ import pytest import tempfile import os -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch from tests.fixtures.databricks.client import DatabricksClientStub from tests.fixtures.amperity import AmperityClientStub @@ -14,6 +14,20 @@ # Import environment fixtures to make them available globally +@pytest.fixture(autouse=True) +def mock_job_cache(): + """ + Automatically mock job cache for all tests to prevent cache pollution. + + This fixture runs automatically for every test and prevents tests from + writing to the user's actual job cache file (~/.chuck_job_cache.json). + Tests that specifically need to test cache behavior should use the + JobCache class directly with a temporary file. + """ + with patch("chuck_data.job_cache.cache_job") as mock_cache: + yield mock_cache + + @pytest.fixture def databricks_client_stub(): """Create a fresh DatabricksClientStub for each test."""