import os from unittest import mock import pytest from click.testing import CliRunner from mlflow.assistant.cli import commands from mlflow.assistant.config import ProviderConfig @pytest.fixture def runner(): return CliRunner() def test_assistant_help(runner): result = runner.invoke(commands, ["--help"]) assert result.exit_code == 0 assert "AI-powered trace analysis" in result.output assert "--configure" in result.output def test_configure_cli_not_found(runner): with mock.patch("mlflow.assistant.providers.claude_code.shutil.which", return_value=None): result = runner.invoke(commands, ["--configure"], input="1\n") assert "not available" in result.output def test_configure_auth_failure(runner): mock_result = mock.Mock() mock_result.returncode = 1 mock_result.stderr = "unauthorized" with ( mock.patch( "mlflow.assistant.providers.claude_code.shutil.which", return_value="/usr/bin/claude" ), mock.patch( "mlflow.assistant.providers.claude_code.subprocess.run", return_value=mock_result, ), ): result = runner.invoke(commands, ["--configure"], input="1\n") assert result.exit_code == 0 assert "not authenticated" in result.output.lower() def test_configure_experiment_fetch_failure(runner): mock_result = mock.Mock() mock_result.returncode = 0 mock_result.stderr = "" with ( mock.patch( "mlflow.assistant.providers.claude_code.shutil.which", return_value="/usr/bin/claude" ), mock.patch( "mlflow.assistant.providers.claude_code.subprocess.run", return_value=mock_result, ), mock.patch( "mlflow.assistant.cli._fetch_recent_experiments", return_value=[], ), ): # Input: provider=1, connect=y, tracking_uri=default result = runner.invoke( commands, ["--configure"], input="1\ny\nhttp://localhost:5000\n", ) assert "Could not fetch experiments" in result.output def test_configure_success(runner, tmp_path): mock_result = mock.Mock() mock_result.returncode = 0 mock_result.stderr = "" mock_config = mock.Mock() mock_config.providers = {"claude_code": ProviderConfig()} mock_config.projects = {} def mock_set_provider(name, model): mock_config.providers[name] = ProviderConfig(model=model, selected=True) mock_config.set_provider = mock_set_provider with ( mock.patch( "mlflow.assistant.providers.claude_code.shutil.which", return_value="/usr/bin/claude" ), mock.patch( "mlflow.assistant.providers.claude_code.subprocess.run", return_value=mock_result, ), mock.patch( "mlflow.assistant.cli._fetch_recent_experiments", return_value=[("1", "Test Experiment")], ), mock.patch( "mlflow.assistant.cli.AssistantConfig.load", return_value=mock_config, ), mock.patch.object(mock_config, "save"), runner.isolated_filesystem(temp_dir=tmp_path), ): # Input: provider=1, connect=y, experiment=1, project_path, model=default, skill_location=1 result = runner.invoke( commands, ["--configure"], input=f"1\ny\nhttp://localhost:5000\n1\n{tmp_path}\ndefault\n1\n", ) assert "Setup Complete" in result.output def test_configure_tilde_expansion(runner): mock_result = mock.Mock() mock_result.returncode = 0 mock_result.stderr = "" mock_config = mock.Mock() mock_config.providers = {"claude_code": ProviderConfig()} projects_dict = {} mock_config.projects = projects_dict def mock_set_provider(name, model): mock_config.providers[name] = ProviderConfig(model=model, selected=True) mock_config.set_provider = mock_set_provider home_dir = os.path.expanduser("~") with ( mock.patch( "mlflow.assistant.providers.claude_code.shutil.which", return_value="/usr/bin/claude" ), mock.patch( "mlflow.assistant.providers.claude_code.subprocess.run", return_value=mock_result, ), mock.patch( "mlflow.assistant.cli._fetch_recent_experiments", return_value=[("1", "Test Experiment")], ), mock.patch( "mlflow.assistant.cli.AssistantConfig.load", return_value=mock_config, ), mock.patch.object(mock_config, "save"), ): # Input: provider=1, connect=y, tracking_uri, experiment=1, project_path=~, # model=default, skill_location=1 result = runner.invoke( commands, ["--configure"], input="1\ny\nhttp://localhost:5000\n1\n~\ndefault\n1\n", ) # Should succeed because ~ expands to home dir which exists assert "Setup Complete" in result.output # Verify the saved path is the expanded path, not ~ assert "1" in projects_dict assert projects_dict["1"].location == home_dir def test_configure_relative_path(runner): mock_result = mock.Mock() mock_result.returncode = 0 mock_result.stderr = "" mock_config = mock.Mock() mock_config.providers = {"claude_code": ProviderConfig()} projects_dict = {} mock_config.projects = projects_dict def mock_set_provider(name, model): mock_config.providers[name] = ProviderConfig(model=model, selected=True) mock_config.set_provider = mock_set_provider with ( mock.patch( "mlflow.assistant.providers.claude_code.shutil.which", return_value="/usr/bin/claude" ), mock.patch( "mlflow.assistant.providers.claude_code.subprocess.run", return_value=mock_result, ), mock.patch( "mlflow.assistant.cli._fetch_recent_experiments", return_value=[("1", "Test Experiment")], ), mock.patch( "mlflow.assistant.cli.AssistantConfig.load", return_value=mock_config, ), mock.patch.object(mock_config, "save"), ): # Use "." which should resolve to current directory # Input: provider=1, connect=y, tracking_uri, experiment=1, project_path=., # model=default, skill_location=1 result = runner.invoke( commands, ["--configure"], input="1\ny\nhttp://localhost:5000\n1\n.\ndefault\n1\n", ) assert "Setup Complete" in result.output # Verify the saved path is absolute, not "." assert "1" in projects_dict assert os.path.isabs(projects_dict["1"].location) assert projects_dict["1"].location != "."