wehub-resource-sync
204 lines
6.5 KiB
Python
204 lines
6.5 KiB
Python
from unittest import mock
|
|
|
|
from click.testing import CliRunner
|
|
|
|
from mlflow.cli import cli
|
|
|
|
|
|
def test_list_commands_cli():
|
|
mock_commands = [
|
|
{
|
|
"key": "genai/analyze_experiment",
|
|
"namespace": "genai",
|
|
"description": "Analyzes an MLflow experiment",
|
|
},
|
|
{
|
|
"key": "ml/train",
|
|
"namespace": "ml",
|
|
"description": "Training helper",
|
|
},
|
|
]
|
|
|
|
with mock.patch("mlflow.ai_commands.list_commands", return_value=mock_commands):
|
|
runner = CliRunner()
|
|
result = runner.invoke(cli, ["ai-commands", "list"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "genai/analyze_experiment: Analyzes an MLflow experiment" in result.output
|
|
assert "ml/train: Training helper" in result.output
|
|
|
|
|
|
def test_list_commands_with_namespace_cli():
|
|
mock_commands = [
|
|
{
|
|
"key": "genai/analyze_experiment",
|
|
"namespace": "genai",
|
|
"description": "Analyzes an MLflow experiment",
|
|
},
|
|
]
|
|
|
|
with mock.patch(
|
|
"mlflow.cli.ai_commands.list_commands", return_value=mock_commands
|
|
) as mock_list:
|
|
runner = CliRunner()
|
|
result = runner.invoke(cli, ["ai-commands", "list", "--namespace", "genai"])
|
|
|
|
assert result.exit_code == 0
|
|
mock_list.assert_called_once_with("genai")
|
|
assert "genai/analyze_experiment" in result.output
|
|
|
|
|
|
def test_list_commands_empty_cli():
|
|
with mock.patch("mlflow.ai_commands.list_commands", return_value=[]):
|
|
runner = CliRunner()
|
|
result = runner.invoke(cli, ["ai-commands", "list"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "No AI commands found" in result.output
|
|
|
|
|
|
def test_list_commands_empty_namespace_cli():
|
|
with mock.patch("mlflow.ai_commands.list_commands", return_value=[]):
|
|
runner = CliRunner()
|
|
result = runner.invoke(cli, ["ai-commands", "list", "--namespace", "unknown"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "No AI commands found in namespace 'unknown'" in result.output
|
|
|
|
|
|
def test_get_command_cli():
|
|
mock_content = """---
|
|
namespace: genai
|
|
description: Test command
|
|
---
|
|
|
|
Hello! This is test content."""
|
|
|
|
with mock.patch("mlflow.ai_commands.get_command", return_value=mock_content):
|
|
runner = CliRunner()
|
|
result = runner.invoke(cli, ["ai-commands", "get", "genai/analyze_experiment"])
|
|
|
|
assert result.exit_code == 0
|
|
assert mock_content == result.output.rstrip("\n")
|
|
|
|
|
|
def test_get_invalid_command_cli():
|
|
with mock.patch(
|
|
"mlflow.cli.ai_commands.get_command",
|
|
side_effect=FileNotFoundError("Command 'invalid/cmd' not found"),
|
|
):
|
|
runner = CliRunner()
|
|
result = runner.invoke(cli, ["ai-commands", "get", "invalid/cmd"])
|
|
|
|
assert result.exit_code != 0
|
|
assert "Error: Command 'invalid/cmd' not found" in result.output
|
|
|
|
|
|
def test_ai_commands_help():
|
|
runner = CliRunner()
|
|
result = runner.invoke(cli, ["ai-commands", "--help"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "Manage MLflow AI commands for LLMs" in result.output
|
|
assert "list" in result.output
|
|
assert "get" in result.output
|
|
assert "run" in result.output
|
|
|
|
|
|
def test_get_command_help():
|
|
runner = CliRunner()
|
|
result = runner.invoke(cli, ["ai-commands", "get", "--help"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "Get a specific AI command by key" in result.output
|
|
assert "KEY" in result.output
|
|
|
|
|
|
def test_list_command_help():
|
|
runner = CliRunner()
|
|
result = runner.invoke(cli, ["ai-commands", "list", "--help"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "List all available AI commands" in result.output
|
|
assert "--namespace" in result.output
|
|
|
|
|
|
def test_run_command_cli():
|
|
mock_content = """---
|
|
namespace: genai
|
|
description: Test command
|
|
---
|
|
|
|
# Test Command
|
|
This is test content."""
|
|
|
|
with mock.patch("mlflow.ai_commands.get_command", return_value=mock_content):
|
|
runner = CliRunner()
|
|
result = runner.invoke(cli, ["ai-commands", "run", "genai/analyze_experiment"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "The user has run an MLflow AI command via CLI" in result.output
|
|
assert "Start executing the workflow immediately without any preamble" in result.output
|
|
assert "# Test Command" in result.output
|
|
assert "This is test content." in result.output
|
|
# Should not have frontmatter
|
|
assert "namespace: genai" not in result.output
|
|
assert "description: Test command" not in result.output
|
|
assert "---" not in result.output
|
|
|
|
|
|
def test_run_invalid_command_cli():
|
|
with mock.patch(
|
|
"mlflow.ai_commands.get_command",
|
|
side_effect=FileNotFoundError("Command 'invalid/cmd' not found"),
|
|
):
|
|
runner = CliRunner()
|
|
result = runner.invoke(cli, ["ai-commands", "run", "invalid/cmd"])
|
|
|
|
assert result.exit_code != 0
|
|
assert "Error: Command 'invalid/cmd' not found" in result.output
|
|
|
|
|
|
def test_run_command_help():
|
|
runner = CliRunner()
|
|
result = runner.invoke(cli, ["ai-commands", "run", "--help"])
|
|
|
|
assert result.exit_code == 0
|
|
assert "Get a command formatted for execution by an AI assistant" in result.output
|
|
assert "KEY" in result.output
|
|
|
|
|
|
def test_actual_command_exists():
|
|
runner = CliRunner()
|
|
|
|
# Test list includes our command
|
|
result = runner.invoke(cli, ["ai-commands", "list"])
|
|
assert result.exit_code == 0
|
|
assert "genai/analyze_experiment" in result.output
|
|
|
|
# Test we can get the command
|
|
result = runner.invoke(cli, ["ai-commands", "get", "genai/analyze_experiment"])
|
|
assert result.exit_code == 0
|
|
assert "# Analyze Experiment" in result.output
|
|
assert "Analyzes traces in an MLflow experiment" in result.output
|
|
|
|
# Test we can run the command
|
|
result = runner.invoke(cli, ["ai-commands", "run", "genai/analyze_experiment"])
|
|
assert result.exit_code == 0
|
|
assert "The user has run an MLflow AI command via CLI" in result.output
|
|
assert "Start executing the workflow immediately without any preamble" in result.output
|
|
assert "# Analyze Experiment" in result.output
|
|
# Should not have frontmatter
|
|
assert "namespace: genai" not in result.output
|
|
assert "---" not in result.output
|
|
|
|
# Test filtering by namespace
|
|
result = runner.invoke(cli, ["ai-commands", "list", "--namespace", "genai"])
|
|
assert result.exit_code == 0
|
|
assert "genai/analyze_experiment" in result.output
|
|
|
|
# Test filtering by wrong namespace excludes it
|
|
result = runner.invoke(cli, ["ai-commands", "list", "--namespace", "ml"])
|
|
assert result.exit_code == 0
|
|
assert "genai/analyze_experiment" not in result.output
|