test_ai_commands.py
1 from unittest import mock 2 3 from click.testing import CliRunner 4 5 from mlflow.cli import cli 6 7 8 def test_list_commands_cli(): 9 mock_commands = [ 10 { 11 "key": "genai/analyze_experiment", 12 "namespace": "genai", 13 "description": "Analyzes an MLflow experiment", 14 }, 15 { 16 "key": "ml/train", 17 "namespace": "ml", 18 "description": "Training helper", 19 }, 20 ] 21 22 with mock.patch("mlflow.ai_commands.list_commands", return_value=mock_commands): 23 runner = CliRunner() 24 result = runner.invoke(cli, ["ai-commands", "list"]) 25 26 assert result.exit_code == 0 27 assert "genai/analyze_experiment: Analyzes an MLflow experiment" in result.output 28 assert "ml/train: Training helper" in result.output 29 30 31 def test_list_commands_with_namespace_cli(): 32 mock_commands = [ 33 { 34 "key": "genai/analyze_experiment", 35 "namespace": "genai", 36 "description": "Analyzes an MLflow experiment", 37 }, 38 ] 39 40 with mock.patch( 41 "mlflow.cli.ai_commands.list_commands", return_value=mock_commands 42 ) as mock_list: 43 runner = CliRunner() 44 result = runner.invoke(cli, ["ai-commands", "list", "--namespace", "genai"]) 45 46 assert result.exit_code == 0 47 mock_list.assert_called_once_with("genai") 48 assert "genai/analyze_experiment" in result.output 49 50 51 def test_list_commands_empty_cli(): 52 with mock.patch("mlflow.ai_commands.list_commands", return_value=[]): 53 runner = CliRunner() 54 result = runner.invoke(cli, ["ai-commands", "list"]) 55 56 assert result.exit_code == 0 57 assert "No AI commands found" in result.output 58 59 60 def test_list_commands_empty_namespace_cli(): 61 with mock.patch("mlflow.ai_commands.list_commands", return_value=[]): 62 runner = CliRunner() 63 result = runner.invoke(cli, ["ai-commands", "list", "--namespace", "unknown"]) 64 65 assert result.exit_code == 0 66 assert "No AI commands found in namespace 'unknown'" in result.output 67 68 69 def test_get_command_cli(): 70 mock_content = """--- 71 namespace: genai 72 description: Test command 73 --- 74 75 Hello! This is test content.""" 76 77 with mock.patch("mlflow.ai_commands.get_command", return_value=mock_content): 78 runner = CliRunner() 79 result = runner.invoke(cli, ["ai-commands", "get", "genai/analyze_experiment"]) 80 81 assert result.exit_code == 0 82 assert mock_content == result.output.rstrip("\n") 83 84 85 def test_get_invalid_command_cli(): 86 with mock.patch( 87 "mlflow.cli.ai_commands.get_command", 88 side_effect=FileNotFoundError("Command 'invalid/cmd' not found"), 89 ): 90 runner = CliRunner() 91 result = runner.invoke(cli, ["ai-commands", "get", "invalid/cmd"]) 92 93 assert result.exit_code != 0 94 assert "Error: Command 'invalid/cmd' not found" in result.output 95 96 97 def test_ai_commands_help(): 98 runner = CliRunner() 99 result = runner.invoke(cli, ["ai-commands", "--help"]) 100 101 assert result.exit_code == 0 102 assert "Manage MLflow AI commands for LLMs" in result.output 103 assert "list" in result.output 104 assert "get" in result.output 105 assert "run" in result.output 106 107 108 def test_get_command_help(): 109 runner = CliRunner() 110 result = runner.invoke(cli, ["ai-commands", "get", "--help"]) 111 112 assert result.exit_code == 0 113 assert "Get a specific AI command by key" in result.output 114 assert "KEY" in result.output 115 116 117 def test_list_command_help(): 118 runner = CliRunner() 119 result = runner.invoke(cli, ["ai-commands", "list", "--help"]) 120 121 assert result.exit_code == 0 122 assert "List all available AI commands" in result.output 123 assert "--namespace" in result.output 124 125 126 def test_run_command_cli(): 127 mock_content = """--- 128 namespace: genai 129 description: Test command 130 --- 131 132 # Test Command 133 This is test content.""" 134 135 with mock.patch("mlflow.ai_commands.get_command", return_value=mock_content): 136 runner = CliRunner() 137 result = runner.invoke(cli, ["ai-commands", "run", "genai/analyze_experiment"]) 138 139 assert result.exit_code == 0 140 assert "The user has run an MLflow AI command via CLI" in result.output 141 assert "Start executing the workflow immediately without any preamble" in result.output 142 assert "# Test Command" in result.output 143 assert "This is test content." in result.output 144 # Should not have frontmatter 145 assert "namespace: genai" not in result.output 146 assert "description: Test command" not in result.output 147 assert "---" not in result.output 148 149 150 def test_run_invalid_command_cli(): 151 with mock.patch( 152 "mlflow.ai_commands.get_command", 153 side_effect=FileNotFoundError("Command 'invalid/cmd' not found"), 154 ): 155 runner = CliRunner() 156 result = runner.invoke(cli, ["ai-commands", "run", "invalid/cmd"]) 157 158 assert result.exit_code != 0 159 assert "Error: Command 'invalid/cmd' not found" in result.output 160 161 162 def test_run_command_help(): 163 runner = CliRunner() 164 result = runner.invoke(cli, ["ai-commands", "run", "--help"]) 165 166 assert result.exit_code == 0 167 assert "Get a command formatted for execution by an AI assistant" in result.output 168 assert "KEY" in result.output 169 170 171 def test_actual_command_exists(): 172 runner = CliRunner() 173 174 # Test list includes our command 175 result = runner.invoke(cli, ["ai-commands", "list"]) 176 assert result.exit_code == 0 177 assert "genai/analyze_experiment" in result.output 178 179 # Test we can get the command 180 result = runner.invoke(cli, ["ai-commands", "get", "genai/analyze_experiment"]) 181 assert result.exit_code == 0 182 assert "# Analyze Experiment" in result.output 183 assert "Analyzes traces in an MLflow experiment" in result.output 184 185 # Test we can run the command 186 result = runner.invoke(cli, ["ai-commands", "run", "genai/analyze_experiment"]) 187 assert result.exit_code == 0 188 assert "The user has run an MLflow AI command via CLI" in result.output 189 assert "Start executing the workflow immediately without any preamble" in result.output 190 assert "# Analyze Experiment" in result.output 191 # Should not have frontmatter 192 assert "namespace: genai" not in result.output 193 assert "---" not in result.output 194 195 # Test filtering by namespace 196 result = runner.invoke(cli, ["ai-commands", "list", "--namespace", "genai"]) 197 assert result.exit_code == 0 198 assert "genai/analyze_experiment" in result.output 199 200 # Test filtering by wrong namespace excludes it 201 result = runner.invoke(cli, ["ai-commands", "list", "--namespace", "ml"]) 202 assert result.exit_code == 0 203 assert "genai/analyze_experiment" not in result.output