1  import json
  2  from pathlib import Path
  3  from unittest import mock
  4  
  5  import pytest
  6  from click.testing import CliRunner
  7  
  8  from mlflow.claude_code.cli import commands
  9  from mlflow.claude_code.config import HOOK_FIELD_COMMAND, HOOK_FIELD_HOOKS
 10  from mlflow.claude_code.hooks import upsert_hook
 11  
 12  
 13  @pytest.fixture
 14  def runner():
 15      """Provide a CLI runner for tests."""
 16      return CliRunner()
 17  
 18  
 19  def test_claude_help_command(runner):
 20      result = runner.invoke(commands, ["--help"])
 21      assert result.exit_code == 0
 22      assert "Commands for autologging with MLflow" in result.output
 23      assert "claude" in result.output
 24  
 25  
 26  def test_trace_command_help(runner):
 27      result = runner.invoke(commands, ["claude", "--help"])
 28      assert result.exit_code == 0
 29      assert "Set up Claude Code tracing" in result.output
 30      assert "--tracking-uri" in result.output
 31      assert "--experiment-id" in result.output
 32      assert "--disable" in result.output
 33      assert "--status" in result.output
 34  
 35  
 36  def test_trace_status_with_no_config(runner):
 37      with runner.isolated_filesystem():
 38          result = runner.invoke(commands, ["claude", "--status"])
 39          assert result.exit_code == 0
 40          assert "❌ Claude tracing is not enabled" in result.output
 41  
 42  
 43  def test_trace_disable_with_no_config(runner):
 44      with runner.isolated_filesystem():
 45          result = runner.invoke(commands, ["claude", "--disable"])
 46          assert result.exit_code == 0
 47  
 48  
 49  def _get_hook_command_from_settings() -> str:
 50      settings_path = Path(".claude/settings.json")
 51      with open(settings_path) as f:
 52          config = json.load(f)
 53  
 54      if hooks := config.get("hooks"):
 55          for group in hooks.get("Stop", []):
 56              for hook in group.get("hooks", []):
 57                  if command := hook.get("command"):
 58                      return command
 59  
 60      raise AssertionError("No hook command found in settings.json")
 61  
 62  
 63  def test_claude_setup_with_uv_env_var(runner, monkeypatch):
 64      monkeypatch.setenv("UV", "/path/to/uv")
 65  
 66      with runner.isolated_filesystem():
 67          result = runner.invoke(commands, ["claude"])
 68          assert result.exit_code == 0
 69  
 70          hook_command = _get_hook_command_from_settings()
 71          assert hook_command == "uv run mlflow autolog claude stop-hook"
 72  
 73  
 74  def test_claude_setup_without_uv_env_var(runner, monkeypatch):
 75      monkeypatch.delenv("UV", raising=False)
 76  
 77      with runner.isolated_filesystem():
 78          result = runner.invoke(commands, ["claude"])
 79          assert result.exit_code == 0
 80  
 81          hook_command = _get_hook_command_from_settings()
 82          assert hook_command == "mlflow autolog claude stop-hook"
 83  
 84  
 85  def test_upsert_hook_uses_cli_command():
 86      config = {HOOK_FIELD_HOOKS: {}}
 87      upsert_hook(config, "Stop", "stop-hook")
 88  
 89      hook_command = config[HOOK_FIELD_HOOKS]["Stop"][0][HOOK_FIELD_HOOKS][0][HOOK_FIELD_COMMAND]
 90      assert "mlflow autolog claude stop-hook" in hook_command
 91  
 92  
 93  def test_upsert_hook_upgrades_legacy_hook():
 94      legacy_command = (
 95          'python -I -c "from mlflow.claude_code.hooks import stop_hook_handler; stop_hook_handler()"'
 96      )
 97      config = {
 98          HOOK_FIELD_HOOKS: {
 99              "Stop": [{HOOK_FIELD_HOOKS: [{"type": "command", HOOK_FIELD_COMMAND: legacy_command}]}]
100          }
101      }
102      upsert_hook(config, "Stop", "stop-hook")
103  
104      hook_command = config[HOOK_FIELD_HOOKS]["Stop"][0][HOOK_FIELD_HOOKS][0][HOOK_FIELD_COMMAND]
105      assert "mlflow autolog claude stop-hook" in hook_command
106      assert "python -I -c" not in hook_command
107  
108  
109  def test_stop_hook_subcommand_is_routable(runner):
110      with mock.patch("mlflow.claude_code.cli.stop_hook_handler") as mock_handler:
111          result = runner.invoke(commands, ["claude", "stop-hook"])
112          assert result.exit_code == 0
113          mock_handler.assert_called_once()