chore: import upstream snapshot with attribution

tests/cli/test_ai_commands.py

@@ -0,0 +1,203 @@
+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

tests/cli/test_crypto.py

@@ -0,0 +1,414 @@
+import logging
+from contextlib import contextmanager
+from unittest import mock
+
+import pytest
+from click.testing import CliRunner
+
+from mlflow.cli.crypto import commands
+from mlflow.exceptions import MlflowException
+
+
+@pytest.fixture(autouse=True)
+def suppress_logging():
+    original_root = logging.root.level
+    original_mlflow = logging.getLogger("mlflow").level
+    original_alembic = logging.getLogger("alembic").level
+
+    logging.root.setLevel(logging.CRITICAL)
+    logging.getLogger("mlflow").setLevel(logging.CRITICAL)
+    logging.getLogger("alembic").setLevel(logging.CRITICAL)
+
+    yield
+
+    logging.root.setLevel(original_root)
+    logging.getLogger("mlflow").setLevel(original_mlflow)
+    logging.getLogger("alembic").setLevel(original_alembic)
+
+
+@pytest.fixture
+def runner():
+    return CliRunner()
+
+
+@pytest.fixture
+def old_passphrase_env(monkeypatch):
+    monkeypatch.setenv("MLFLOW_CRYPTO_KEK_PASSPHRASE", "old-passphrase")
+
+
+@pytest.fixture
+def mock_session():
+    session = mock.MagicMock()
+    session.__enter__ = mock.Mock(return_value=session)
+    session.__exit__ = mock.Mock(return_value=False)
+    return session
+
+
+@pytest.fixture
+def mock_store(mock_session):
+    store = mock.Mock()
+    store.ManagedSessionMaker.return_value = mock_session
+    return store
+
+
+@pytest.fixture
+def mock_secret():
+    secret = mock.Mock()
+    secret.secret_id = "test-secret-id-123"
+    secret.encrypted_value = b"encrypted-data"
+    secret.wrapped_dek = b"wrapped-dek-data"
+    secret.kek_version = 1
+    return secret
+
+
+@pytest.fixture
+def empty_db(mock_session):
+    mock_session.query.return_value.filter.return_value.all.return_value = []
+
+
+@pytest.fixture
+def db_with_secret(mock_session, mock_secret):
+    mock_session.query.return_value.filter.return_value.all.return_value = [mock_secret]
+
+
+@contextmanager
+def patch_backend(mock_store):
+    mock_sql_secret = mock.Mock()
+    with (
+        mock.patch("mlflow.cli.crypto._get_store", return_value=mock_store),
+        mock.patch.dict(
+            "sys.modules",
+            {"mlflow.store.tracking.dbmodels.models": mock.Mock(SqlGatewaySecret=mock_sql_secret)},
+        ),
+    ):
+        yield mock_sql_secret
+
+
+@contextmanager
+def patch_rotation(return_value=None):
+    result = mock.Mock()
+    result.wrapped_dek = b"new-wrapped-dek"
+    with mock.patch(
+        "mlflow.cli.crypto.rotate_secret_encryption",
+        return_value=return_value or result,
+    ):
+        yield
+
+
+def test_crypto_group_exists():
+    assert commands.name == "crypto"
+    assert commands.help is not None
+    assert "cryptographic" in commands.help.lower()
+
+
+def test_rotate_kek_command_exists():
+    rotate_cmd = next((cmd for cmd in commands.commands.values() if cmd.name == "rotate-kek"), None)
+    assert rotate_cmd is not None
+    assert "rotate" in rotate_cmd.help.lower()
+
+
+def test_rotate_kek_has_required_parameters():
+    rotate_cmd = next((cmd for cmd in commands.commands.values() if cmd.name == "rotate-kek"), None)
+    param_names = [p.name for p in rotate_cmd.params]
+    assert "new_passphrase" in param_names
+    assert "backend_store_uri" in param_names
+    assert "yes" in param_names
+
+
+def test_new_passphrase_is_required():
+    rotate_cmd = next((cmd for cmd in commands.commands.values() if cmd.name == "rotate-kek"), None)
+    new_pass_param = next((p for p in rotate_cmd.params if p.name == "new_passphrase"), None)
+    assert new_pass_param.required
+    assert new_pass_param.prompt
+    assert new_pass_param.hide_input
+    assert new_pass_param.confirmation_prompt
+
+
+def test_yes_flag_is_optional():
+    rotate_cmd = next((cmd for cmd in commands.commands.values() if cmd.name == "rotate-kek"), None)
+    yes_param = next((p for p in rotate_cmd.params if p.name == "yes"), None)
+    assert yes_param.is_flag
+    assert yes_param.default is False
+
+
+def test_missing_old_passphrase_raises_error(runner, monkeypatch):
+    monkeypatch.delenv("MLFLOW_CRYPTO_KEK_PASSPHRASE", raising=False)
+    result = runner.invoke(commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"])
+    assert result.exit_code != 0
+    assert result.exception is not None
+    assert "MLFLOW_CRYPTO_KEK_PASSPHRASE" in str(result.exception)
+
+
+def test_old_passphrase_from_env(runner, old_passphrase_env, mock_store, empty_db):
+    with patch_backend(mock_store):
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"]
+        )
+        assert result.exit_code == 0
+        assert "No secrets found" in result.output
+
+
+def test_kek_version_defaults_to_1(runner, old_passphrase_env, mock_store, empty_db, monkeypatch):
+    monkeypatch.delenv("MLFLOW_CRYPTO_KEK_VERSION", raising=False)
+    with patch_backend(mock_store), mock.patch("mlflow.cli.crypto.KEKManager") as mock_kek:
+        runner.invoke(commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"])
+        assert mock_kek.call_args_list[0][1]["kek_version"] == 1
+
+
+def test_kek_version_from_env(runner, mock_store, empty_db, monkeypatch):
+    monkeypatch.setenv("MLFLOW_CRYPTO_KEK_PASSPHRASE", "old-passphrase")
+    monkeypatch.setenv("MLFLOW_CRYPTO_KEK_VERSION", "5")
+    with patch_backend(mock_store):
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"]
+        )
+        assert result.exit_code == 0
+
+
+def test_version_increments_correctly(runner, mock_store, empty_db, monkeypatch):
+    monkeypatch.setenv("MLFLOW_CRYPTO_KEK_PASSPHRASE", "old-passphrase")
+    monkeypatch.setenv("MLFLOW_CRYPTO_KEK_VERSION", "3")
+    with patch_backend(mock_store):
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"]
+        )
+        assert result.exit_code == 0
+
+
+def test_interactive_prompt_shows_warning(runner, old_passphrase_env, mock_store, empty_db):
+    with patch_backend(mock_store):
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase"], input="n\n"
+        )
+        assert "⚠️  WARNING: KEK Rotation Operation" in result.output
+        assert "Re-wrap all encryption DEKs" in result.output
+        assert "MLFLOW_CRYPTO_KEK_PASSPHRASE" in result.output
+        assert "MLFLOW_CRYPTO_KEK_VERSION" in result.output
+        assert "Continue with KEK rotation?" in result.output
+
+
+def test_yes_flag_skips_confirmation(runner, old_passphrase_env, mock_store, empty_db):
+    with patch_backend(mock_store):
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"]
+        )
+        assert "⚠️  WARNING" not in result.output
+        assert "Continue with KEK rotation?" not in result.output
+
+
+def test_cancellation_exits_gracefully(runner, old_passphrase_env, mock_store, empty_db):
+    with patch_backend(mock_store):
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase"], input="n\n"
+        )
+        assert result.exit_code == 0
+        assert "KEK rotation cancelled" in result.output
+
+
+def test_connects_to_backend_store(runner, old_passphrase_env, mock_store, empty_db):
+    with patch_backend(mock_store):
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"]
+        )
+        assert result.exit_code == 0
+
+
+def test_uses_custom_backend_store_uri(runner, old_passphrase_env, mock_store, empty_db):
+    with patch_backend(mock_store):
+        result = runner.invoke(
+            commands,
+            [
+                "rotate-kek",
+                "--new-passphrase",
+                "new-passphrase",
+                "--backend-store-uri",
+                "sqlite:///test.db",
+                "--yes",
+            ],
+        )
+        assert result.exit_code == 0
+
+
+def test_filters_secrets_by_kek_version(
+    runner, old_passphrase_env, mock_store, mock_secret, monkeypatch
+):
+    monkeypatch.setenv("MLFLOW_CRYPTO_KEK_VERSION", "2")
+    mock_session = mock_store.ManagedSessionMaker.return_value
+    mock_session.query.return_value.filter.return_value.all.return_value = [mock_secret]
+    with patch_backend(mock_store), patch_rotation():
+        runner.invoke(commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"])
+        assert mock_session.query.return_value.filter.call_args is not None
+
+
+def test_commits_transaction_on_success(runner, old_passphrase_env, mock_store, db_with_secret):
+    with patch_backend(mock_store), patch_rotation():
+        runner.invoke(commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"])
+        mock_store.ManagedSessionMaker.return_value.commit.assert_called_once()
+
+
+def test_no_secrets_returns_success(runner, old_passphrase_env, mock_store, empty_db):
+    with patch_backend(mock_store):
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"]
+        )
+        assert result.exit_code == 0
+        assert "No secrets found" in result.output
+        assert "Nothing to rotate" in result.output
+
+
+def test_wrong_old_passphrase_fails(runner, old_passphrase_env, mock_store, db_with_secret):
+    with (
+        patch_backend(mock_store),
+        mock.patch(
+            "mlflow.cli.crypto.rotate_secret_encryption",
+            side_effect=MlflowException("Failed to rotate secret encryption"),
+        ),
+    ):
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"]
+        )
+        assert result.exit_code != 0
+        assert "Failed to rotate secret" in result.output
+
+
+def test_rotation_failure_rolls_back(runner, old_passphrase_env, mock_store, db_with_secret):
+    with (
+        patch_backend(mock_store),
+        mock.patch(
+            "mlflow.cli.crypto.rotate_secret_encryption",
+            side_effect=Exception("Rotation failed"),
+        ),
+    ):
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"]
+        )
+        assert result.exit_code != 0
+        mock_store.ManagedSessionMaker.return_value.rollback.assert_called_once()
+        assert "No changes were made" in str(result.exception)
+
+
+def test_database_connection_error(runner, old_passphrase_env):
+    mock_sql_secret = mock.Mock()
+    with (
+        mock.patch("mlflow.cli.crypto._get_store", side_effect=Exception("Connection failed")),
+        mock.patch.dict(
+            "sys.modules",
+            {"mlflow.store.tracking.dbmodels.models": mock.Mock(SqlGatewaySecret=mock_sql_secret)},
+        ),
+    ):
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"]
+        )
+        assert result.exit_code != 0
+        assert "Failed to connect to backend store" in str(result.exception)
+
+
+def test_kek_manager_creation_error(runner, old_passphrase_env, mock_store):
+    with (
+        patch_backend(mock_store),
+        mock.patch("mlflow.cli.crypto.KEKManager", side_effect=Exception("KEK creation failed")),
+    ):
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"]
+        )
+        assert result.exit_code != 0
+        assert "Failed to create KEK managers" in str(result.exception)
+
+
+def test_shows_progress_for_multiple_secrets(runner, old_passphrase_env, mock_store, mock_session):
+    secrets = [
+        mock.Mock(
+            secret_id=f"secret-{i}", encrypted_value=b"enc", wrapped_dek=b"wrap", kek_version=1
+        )
+        for i in range(5)
+    ]
+    mock_session.query.return_value.filter.return_value.all.return_value = secrets
+    with patch_backend(mock_store), patch_rotation():
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"]
+        )
+        assert result.exit_code == 0
+        assert "Found 5 secrets to rotate" in result.output
+
+
+def test_success_message_includes_version_info(runner, mock_store, mock_session, monkeypatch):
+    monkeypatch.setenv("MLFLOW_CRYPTO_KEK_PASSPHRASE", "old-passphrase")
+    monkeypatch.setenv("MLFLOW_CRYPTO_KEK_VERSION", "3")
+    secret = mock.Mock(secret_id="test", encrypted_value=b"enc", wrapped_dek=b"wrap", kek_version=3)
+    mock_session.query.return_value.filter.return_value.all.return_value = [secret]
+    with patch_backend(mock_store), patch_rotation():
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"]
+        )
+        assert result.exit_code == 0
+        assert "Successfully rotated 1 encryption key" in result.output
+
+
+def test_shows_environment_variable_warning(runner, old_passphrase_env, mock_store, db_with_secret):
+    with patch_backend(mock_store), patch_rotation():
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"]
+        )
+        assert result.exit_code == 0
+        assert "CRITICAL: Update BOTH environment variables" in result.output
+        assert "MLFLOW_CRYPTO_KEK_PASSPHRASE='<new-passphrase>'" in result.output
+        assert "MLFLOW_CRYPTO_KEK_VERSION='2'" in result.output
+        assert "Failure to update BOTH variables will cause decryption failures" in result.output
+
+
+def test_large_number_of_secrets(runner, old_passphrase_env, mock_store, mock_session):
+    secrets = [
+        mock.Mock(
+            secret_id=f"secret-{i}", encrypted_value=b"enc", wrapped_dek=b"wrap", kek_version=1
+        )
+        for i in range(1000)
+    ]
+    mock_session.query.return_value.filter.return_value.all.return_value = secrets
+    with patch_backend(mock_store), patch_rotation():
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"]
+        )
+        assert result.exit_code == 0
+        assert "Found 1000 secrets to rotate" in result.output
+        assert "Successfully rotated 1000 encryption keys" in result.output
+
+
+def test_mixed_kek_versions_only_rotates_old_version(
+    runner, old_passphrase_env, mock_store, mock_session, monkeypatch
+):
+    monkeypatch.setenv("MLFLOW_CRYPTO_KEK_VERSION", "1")
+    secret = mock.Mock(
+        secret_id="secret-v1", encrypted_value=b"enc", wrapped_dek=b"wrap", kek_version=1
+    )
+    mock_session.query.return_value.filter.return_value.all.return_value = [secret]
+    with patch_backend(mock_store), patch_rotation():
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"]
+        )
+        assert result.exit_code == 0
+        assert "Found 1 secrets to rotate" in result.output
+
+
+def test_rotation_with_special_characters_in_passphrase(runner, mock_store, empty_db, monkeypatch):
+    monkeypatch.setenv("MLFLOW_CRYPTO_KEK_PASSPHRASE", "old-p@$$phrase!#$%")
+    with patch_backend(mock_store):
+        result = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-p@$$phrase!#$%", "--yes"]
+        )
+        assert result.exit_code == 0
+
+
+def test_idempotency_running_twice_with_same_version(
+    runner, old_passphrase_env, mock_store, empty_db
+):
+    with patch_backend(mock_store):
+        result1 = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase", "--yes"]
+        )
+        result2 = runner.invoke(
+            commands, ["rotate-kek", "--new-passphrase", "new-passphrase-2", "--yes"]
+        )
+        assert result1.exit_code == 0
+        assert result2.exit_code == 0
+        assert "No secrets found" in result1.output
+        assert "No secrets found" in result2.output

tests/cli/test_datasets.py

@@ -0,0 +1,184 @@
+import json
+
+import pytest
+from click.testing import CliRunner
+
+import mlflow
+from mlflow.cli.datasets import commands
+from mlflow.genai.datasets import create_dataset
+
+
+@pytest.fixture
+def runner():
+    return CliRunner(catch_exceptions=False)
+
+
+@pytest.fixture
+def experiment():
+    exp_id = mlflow.create_experiment("test_datasets_cli")
+    yield exp_id
+    mlflow.delete_experiment(exp_id)
+
+
+@pytest.fixture
+def dataset_a(experiment):
+    return create_dataset(
+        name="dataset_a",
+        experiment_id=experiment,
+        tags={"env": "production"},
+    )
+
+
+@pytest.fixture
+def dataset_b(experiment):
+    return create_dataset(
+        name="dataset_b",
+        experiment_id=experiment,
+        tags={"env": "staging"},
+    )
+
+
+def test_commands_group_exists():
+    assert commands.name == "datasets"
+    assert commands.help is not None
+
+
+def test_list_command_params():
+    list_cmd = next((cmd for cmd in commands.commands.values() if cmd.name == "list"), None)
+    assert list_cmd is not None
+    param_names = {p.name for p in list_cmd.params}
+    expected_params = {
+        "experiment_id",
+        "filter_string",
+        "max_results",
+        "order_by",
+        "page_token",
+        "output",
+    }
+    assert param_names == expected_params
+
+
+def test_list_datasets_table_output(runner: CliRunner, experiment: str, dataset_a):
+    result = runner.invoke(commands, ["list", "--experiment-id", experiment])
+
+    assert result.exit_code == 0
+    assert dataset_a.dataset_id in result.output
+    assert "dataset_a" in result.output
+
+
+def test_list_datasets_json_output(runner: CliRunner, experiment: str, dataset_a):
+    result = runner.invoke(commands, ["list", "--experiment-id", experiment, "--output", "json"])
+
+    assert result.exit_code == 0
+
+    expected = {
+        "datasets": [
+            {
+                "dataset_id": dataset_a.dataset_id,
+                "name": "dataset_a",
+                "digest": dataset_a.digest,
+                "created_time": dataset_a.created_time,
+                "last_update_time": dataset_a.last_update_time,
+                "created_by": dataset_a.created_by,
+                "last_updated_by": dataset_a.last_updated_by,
+                "tags": dataset_a.tags,
+            }
+        ],
+        "next_page_token": None,
+    }
+    assert json.loads(result.output) == expected
+
+
+def test_list_datasets_empty_results(runner: CliRunner, experiment: str):
+    result = runner.invoke(commands, ["list", "--experiment-id", experiment])
+
+    assert result.exit_code == 0
+
+
+def test_list_datasets_json_empty_results(runner: CliRunner, experiment: str):
+    result = runner.invoke(commands, ["list", "--experiment-id", experiment, "--output", "json"])
+
+    assert result.exit_code == 0
+    output_json = json.loads(result.output)
+    assert output_json == {"datasets": [], "next_page_token": None}
+
+
+def test_list_datasets_with_experiment_id_env_var(runner: CliRunner, experiment: str, dataset_a):
+    result = runner.invoke(commands, ["list"], env={"MLFLOW_EXPERIMENT_ID": experiment})
+
+    assert result.exit_code == 0
+    assert dataset_a.dataset_id in result.output
+
+
+def test_list_datasets_missing_experiment_id(runner: CliRunner):
+    result = runner.invoke(commands, ["list"])
+
+    assert result.exit_code != 0
+    assert "Missing option '--experiment-id' / '-x'" in result.output
+
+
+def test_list_datasets_invalid_output_format(runner: CliRunner, experiment: str):
+    result = runner.invoke(commands, ["list", "--experiment-id", experiment, "--output", "invalid"])
+
+    assert result.exit_code != 0
+    assert "'invalid' is not one of 'table', 'json'" in result.output
+
+
+def test_list_datasets_with_filter_string(runner: CliRunner, experiment: str, dataset_a, dataset_b):
+    result = runner.invoke(
+        commands,
+        ["list", "--experiment-id", experiment, "--filter-string", "name = 'dataset_a'"],
+    )
+
+    assert result.exit_code == 0
+    assert "dataset_a" in result.output
+    assert "dataset_b" not in result.output
+
+
+def test_list_datasets_with_max_results(runner: CliRunner, experiment: str, dataset_a, dataset_b):
+    result = runner.invoke(commands, ["list", "--experiment-id", experiment, "--max-results", "1"])
+
+    assert result.exit_code == 0
+    output_lines = [line for line in result.output.split("\n") if "dataset_" in line]
+    assert len(output_lines) == 1
+
+
+def test_list_datasets_with_order_by(runner: CliRunner, experiment: str, dataset_a, dataset_b):
+    result = runner.invoke(
+        commands, ["list", "--experiment-id", experiment, "--order-by", "name ASC"]
+    )
+
+    assert result.exit_code == 0
+    a_pos = result.output.find("dataset_a")
+    b_pos = result.output.find("dataset_b")
+    assert a_pos < b_pos
+
+
+def test_list_datasets_short_option_x(runner: CliRunner, experiment: str, dataset_a):
+    result = runner.invoke(commands, ["list", "-x", experiment])
+
+    assert result.exit_code == 0
+    assert dataset_a.dataset_id in result.output
+
+
+def test_list_datasets_multiple_datasets(runner: CliRunner, experiment: str, dataset_a, dataset_b):
+    result = runner.invoke(commands, ["list", "--experiment-id", experiment])
+
+    assert result.exit_code == 0
+    assert dataset_a.dataset_id in result.output
+    assert "dataset_a" in result.output
+    assert dataset_b.dataset_id in result.output
+    assert "dataset_b" in result.output
+
+
+def test_list_datasets_json_multiple_datasets(
+    runner: CliRunner, experiment: str, dataset_a, dataset_b
+):
+    result = runner.invoke(commands, ["list", "--experiment-id", experiment, "--output", "json"])
+
+    assert result.exit_code == 0
+    output_json = json.loads(result.output)
+    assert len(output_json["datasets"]) == 2
+
+    names = {d["name"] for d in output_json["datasets"]}
+    assert names == {"dataset_a", "dataset_b"}

tests/cli/test_eval.py

@@ -0,0 +1,217 @@
+import re
+from unittest import mock
+
+import click
+import pandas as pd
+import pytest
+
+import mlflow
+from mlflow.cli.eval import evaluate_traces
+from mlflow.entities import Trace, TraceInfo
+from mlflow.genai.scorers.base import scorer
+
+
+def test_evaluate_traces_with_single_trace_table_output():
+    experiment_id = mlflow.create_experiment("test_experiment")
+
+    mock_trace = mock.Mock(spec=Trace)
+    mock_trace.info = mock.Mock(spec=TraceInfo)
+    mock_trace.info.trace_id = "tr-test-123"
+    mock_trace.info.experiment_id = experiment_id
+
+    mock_results = mock.Mock()
+    mock_results.run_id = "run-eval-456"
+    mock_results.result_df = pd.DataFrame([
+        {
+            "trace_id": "tr-test-123",
+            "assessments": [
+                {
+                    "assessment_name": "RelevanceToQuery",
+                    "feedback": {"value": "yes"},
+                    "rationale": "The answer is relevant",
+                    "metadata": {"mlflow.assessment.sourceRunId": "run-eval-456"},
+                }
+            ],
+        }
+    ])
+
+    with (
+        mock.patch(
+            "mlflow.cli.eval.MlflowClient.get_trace", return_value=mock_trace
+        ) as mock_get_trace,
+        mock.patch("mlflow.cli.eval.evaluate", return_value=mock_results) as mock_evaluate,
+    ):
+        evaluate_traces(
+            experiment_id=experiment_id,
+            trace_ids="tr-test-123",
+            scorers="RelevanceToQuery",
+            output_format="table",
+        )
+
+        mock_get_trace.assert_called_once_with("tr-test-123", display=False)
+
+        assert mock_evaluate.call_count == 1
+        call_args = mock_evaluate.call_args
+        assert "data" in call_args.kwargs
+
+        expected_df = pd.DataFrame([{"trace_id": "tr-test-123", "trace": mock_trace}])
+        pd.testing.assert_frame_equal(call_args.kwargs["data"], expected_df)
+
+        assert "scorers" in call_args.kwargs
+        assert len(call_args.kwargs["scorers"]) == 1
+        assert call_args.kwargs["scorers"][0].__class__.__name__ == "RelevanceToQuery"
+
+
+def test_evaluate_traces_with_multiple_traces_json_output():
+    experiment = mlflow.create_experiment("test_experiment_multi")
+
+    mock_trace1 = mock.Mock(spec=Trace)
+    mock_trace1.info = mock.Mock(spec=TraceInfo)
+    mock_trace1.info.trace_id = "tr-test-1"
+    mock_trace1.info.experiment_id = experiment
+
+    mock_trace2 = mock.Mock(spec=Trace)
+    mock_trace2.info = mock.Mock(spec=TraceInfo)
+    mock_trace2.info.trace_id = "tr-test-2"
+    mock_trace2.info.experiment_id = experiment
+
+    mock_results = mock.Mock()
+    mock_results.run_id = "run-eval-789"
+    mock_results.result_df = pd.DataFrame([
+        {
+            "trace_id": "tr-test-1",
+            "assessments": [
+                {
+                    "assessment_name": "Correctness",
+                    "feedback": {"value": "correct"},
+                    "rationale": "Content is correct",
+                    "metadata": {"mlflow.assessment.sourceRunId": "run-eval-789"},
+                }
+            ],
+        },
+        {
+            "trace_id": "tr-test-2",
+            "assessments": [
+                {
+                    "assessment_name": "Correctness",
+                    "feedback": {"value": "correct"},
+                    "rationale": "Also correct",
+                    "metadata": {"mlflow.assessment.sourceRunId": "run-eval-789"},
+                }
+            ],
+        },
+    ])
+
+    with (
+        mock.patch(
+            "mlflow.cli.eval.MlflowClient.get_trace",
+            side_effect=[mock_trace1, mock_trace2],
+        ) as mock_get_trace,
+        mock.patch("mlflow.cli.eval.evaluate", return_value=mock_results) as mock_evaluate,
+    ):
+        evaluate_traces(
+            experiment_id=experiment,
+            trace_ids="tr-test-1,tr-test-2",
+            scorers="Correctness",
+            output_format="json",
+        )
+
+        assert mock_get_trace.call_count == 2
+        mock_get_trace.assert_any_call("tr-test-1", display=False)
+        mock_get_trace.assert_any_call("tr-test-2", display=False)
+
+        assert mock_evaluate.call_count == 1
+        call_args = mock_evaluate.call_args
+        expected_df = pd.DataFrame([
+            {"trace_id": "tr-test-1", "trace": mock_trace1},
+            {"trace_id": "tr-test-2", "trace": mock_trace2},
+        ])
+        pd.testing.assert_frame_equal(call_args.kwargs["data"], expected_df)
+
+
+def test_evaluate_traces_with_nonexistent_trace():
+    experiment = mlflow.create_experiment("test_experiment_error")
+
+    with mock.patch("mlflow.cli.eval.MlflowClient.get_trace", return_value=None) as mock_get_trace:
+        with pytest.raises(click.UsageError, match="Trace with ID 'tr-nonexistent' not found"):
+            evaluate_traces(
+                experiment_id=experiment,
+                trace_ids="tr-nonexistent",
+                scorers="RelevanceToQuery",
+                output_format="table",
+            )
+
+        mock_get_trace.assert_called_once_with("tr-nonexistent", display=False)
+
+
+def test_evaluate_traces_with_trace_from_wrong_experiment():
+    experiment1 = mlflow.create_experiment("test_experiment_1")
+    experiment2 = mlflow.create_experiment("test_experiment_2")
+
+    mock_trace = mock.Mock(spec=Trace)
+    mock_trace.info = mock.Mock(spec=TraceInfo)
+    mock_trace.info.trace_id = "tr-test-123"
+    mock_trace.info.experiment_id = experiment2
+
+    with mock.patch(
+        "mlflow.cli.eval.MlflowClient.get_trace", return_value=mock_trace
+    ) as mock_get_trace:
+        with pytest.raises(click.UsageError, match="belongs to experiment"):
+            evaluate_traces(
+                experiment_id=experiment1,
+                trace_ids="tr-test-123",
+                scorers="RelevanceToQuery",
+                output_format="table",
+            )
+
+        mock_get_trace.assert_called_once_with("tr-test-123", display=False)
+
+
+def test_evaluate_traces_integration():
+    experiment_id = mlflow.create_experiment("test_experiment_integration")
+    mlflow.set_experiment(experiment_id=experiment_id)
+
+    # Create a few real traces with inputs and outputs
+    trace_ids = []
+    for i in range(3):
+        with mlflow.start_span(name=f"test_span_{i}") as span:
+            span.set_inputs({"question": f"What is test {i}?"})
+            span.set_outputs(f"This is answer {i}")
+            trace_ids.append(span.trace_id)
+
+    # Define a simple code-based scorer inline
+    @scorer
+    def simple_scorer(outputs):
+        """Extract the digit from the output string and return it as the score"""
+        if match := re.search(r"\d+", outputs):
+            return float(match.group())
+        return 0.0
+
+    with mock.patch(
+        "mlflow.cli.eval.resolve_scorers", return_value=[simple_scorer]
+    ) as mock_resolve:
+        evaluate_traces(
+            experiment_id=experiment_id,
+            trace_ids=",".join(trace_ids),
+            scorers="simple_scorer",  # This will be intercepted by our mock
+            output_format="table",
+        )
+        mock_resolve.assert_called_once()
+
+    # Verify that the evaluation results are as expected
+    traces = mlflow.search_traces(locations=[experiment_id], return_type="list")
+    assert len(traces) == 3
+
+    # Sort traces by their outputs to get consistent ordering
+    traces = sorted(traces, key=lambda t: t.data.spans[0].outputs)
+
+    for i, trace in enumerate(traces):
+        assessments = trace.info.assessments
+        assert len(assessments) > 0
+
+        scorer_assessments = [a for a in assessments if a.name == "simple_scorer"]
+        assert len(scorer_assessments) == 1
+
+        assessment = scorer_assessments[0]
+        # Each trace should have a score equal to its index (0, 1, 2)
+        assert assessment.value == float(i)

tests/cli/test_genai_eval_utils.py

@@ -0,0 +1,477 @@
+from unittest import mock
+
+import click
+import pandas as pd
+import pytest
+
+from mlflow.cli.genai_eval_utils import (
+    NA_VALUE,
+    Assessment,
+    EvalResult,
+    extract_assessments_from_results,
+    format_table_output,
+    resolve_scorers,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.tracing.constant import AssessmentMetadataKey
+
+
+def test_format_single_trace_with_result_and_rationale():
+    output_data = [
+        EvalResult(
+            trace_id="tr-123",
+            assessments=[
+                Assessment(
+                    name="RelevanceToQuery",
+                    result="yes",
+                    rationale="The answer is relevant",
+                )
+            ],
+        )
+    ]
+
+    table_output = format_table_output(output_data)
+
+    # Headers should use assessment names from output_data
+    assert table_output.headers == ["trace_id", "RelevanceToQuery"]
+    assert len(table_output.rows) == 1
+    assert table_output.rows[0][0].value == "tr-123"
+    assert "value: yes" in table_output.rows[0][1].value
+    assert "rationale: The answer is relevant" in table_output.rows[0][1].value
+
+
+def test_format_multiple_traces_multiple_scorers():
+    output_data = [
+        EvalResult(
+            trace_id="tr-123",
+            assessments=[
+                Assessment(
+                    name="RelevanceToQuery",
+                    result="yes",
+                    rationale="Relevant",
+                ),
+                Assessment(name="Safety", result="yes", rationale="Safe"),
+            ],
+        ),
+        EvalResult(
+            trace_id="tr-456",
+            assessments=[
+                Assessment(
+                    name="RelevanceToQuery",
+                    result="no",
+                    rationale="Not relevant",
+                ),
+                Assessment(name="Safety", result="yes", rationale="Safe"),
+            ],
+        ),
+    ]
+
+    table_output = format_table_output(output_data)
+
+    # Assessment names should be sorted
+    assert table_output.headers == ["trace_id", "RelevanceToQuery", "Safety"]
+    assert len(table_output.rows) == 2
+    assert table_output.rows[0][0].value == "tr-123"
+    assert table_output.rows[1][0].value == "tr-456"
+    assert "value: yes" in table_output.rows[0][1].value
+    assert "value: no" in table_output.rows[1][1].value
+
+
+def test_format_long_rationale_not_truncated():
+    long_rationale = "x" * 150
+    output_data = [
+        EvalResult(
+            trace_id="tr-123",
+            assessments=[
+                Assessment(
+                    name="RelevanceToQuery",
+                    result="yes",
+                    rationale=long_rationale,
+                )
+            ],
+        )
+    ]
+
+    table_output = format_table_output(output_data)
+
+    assert long_rationale in table_output.rows[0][1].value
+    assert len(table_output.rows[0][1].value) >= len(long_rationale)
+
+
+def test_format_error_message_formatting():
+    output_data = [
+        EvalResult(
+            trace_id="tr-123",
+            assessments=[
+                Assessment(
+                    name="RelevanceToQuery",
+                    result=None,
+                    rationale=None,
+                    error="OpenAI API error",
+                )
+            ],
+        )
+    ]
+
+    table_output = format_table_output(output_data)
+
+    assert table_output.rows[0][1].value == "error: OpenAI API error"
+
+
+def test_format_na_for_missing_results():
+    output_data = [
+        EvalResult(
+            trace_id="tr-123",
+            assessments=[
+                Assessment(
+                    name="RelevanceToQuery",
+                    result=None,
+                    rationale=None,
+                )
+            ],
+        )
+    ]
+
+    table_output = format_table_output(output_data)
+
+    assert table_output.rows[0][1].value == NA_VALUE
+
+
+def test_format_result_only_without_rationale():
+    output_data = [
+        EvalResult(
+            trace_id="tr-123",
+            assessments=[
+                Assessment(
+                    name="RelevanceToQuery",
+                    result="yes",
+                    rationale=None,
+                )
+            ],
+        )
+    ]
+
+    table_output = format_table_output(output_data)
+
+    assert table_output.rows[0][1].value == "value: yes"
+
+
+def test_format_rationale_only_without_result():
+    output_data = [
+        EvalResult(
+            trace_id="tr-123",
+            assessments=[
+                Assessment(
+                    name="RelevanceToQuery",
+                    result=None,
+                    rationale="Some reasoning",
+                )
+            ],
+        )
+    ]
+
+    table_output = format_table_output(output_data)
+
+    assert table_output.rows[0][1].value == "rationale: Some reasoning"
+
+
+def test_format_with_different_assessment_names():
+    # This test demonstrates that assessment names (e.g., "relevance_to_query")
+    # are used in headers, not scorer class names (e.g., "RelevanceToQuery")
+    output_data = [
+        EvalResult(
+            trace_id="tr-123",
+            assessments=[
+                Assessment(
+                    name="relevance_to_query",  # Different from scorer name
+                    result="yes",
+                    rationale="The answer is relevant",
+                ),
+                Assessment(
+                    name="safety_check",  # Different from scorer name
+                    result="safe",
+                    rationale="Content is safe",
+                ),
+            ],
+        )
+    ]
+
+    table_output = format_table_output(output_data)
+
+    # Headers should use actual assessment names from output_data (sorted)
+    assert table_output.headers == ["trace_id", "relevance_to_query", "safety_check"]
+    assert len(table_output.rows) == 1
+    assert table_output.rows[0][0].value == "tr-123"
+    assert "value: yes" in table_output.rows[0][1].value
+    assert "value: safe" in table_output.rows[0][2].value
+
+
+# Tests for resolve_scorers function
+
+
+def test_resolve_builtin_scorer():
+    # Test with real built-in scorer names
+    scorers = resolve_scorers(["Correctness"], "experiment_123")
+
+    assert len(scorers) == 1
+    assert scorers[0].__class__.__name__ == "Correctness"
+
+
+def test_resolve_builtin_scorer_snake_case():
+    # Test with snake_case name
+    scorers = resolve_scorers(["correctness"], "experiment_123")
+
+    assert len(scorers) == 1
+    assert scorers[0].__class__.__name__ == "Correctness"
+
+
+def test_resolve_registered_scorer():
+    mock_registered = mock.Mock()
+
+    with (
+        mock.patch(
+            "mlflow.cli.genai_eval_utils.get_all_scorers", return_value=[]
+        ) as mock_get_all_scorers,
+        mock.patch(
+            "mlflow.cli.genai_eval_utils.get_scorer", return_value=mock_registered
+        ) as mock_get_scorer,
+    ):
+        scorers = resolve_scorers(["CustomScorer"], "experiment_123")
+
+        assert len(scorers) == 1
+        assert scorers[0] == mock_registered
+        # Verify mocks were called as expected
+        mock_get_all_scorers.assert_called_once()
+        mock_get_scorer.assert_called_once_with(name="CustomScorer", experiment_id="experiment_123")
+
+
+def test_resolve_mixed_scorers():
+    # Setup built-in scorer
+    mock_builtin = mock.Mock()
+    mock_builtin.__class__.__name__ = "Safety"
+    mock_builtin.name = None
+
+    # Setup registered scorer
+    mock_registered = mock.Mock()
+
+    with (
+        mock.patch(
+            "mlflow.cli.genai_eval_utils.get_all_scorers", return_value=[mock_builtin]
+        ) as mock_get_all_scorers,
+        mock.patch(
+            "mlflow.cli.genai_eval_utils.get_scorer", return_value=mock_registered
+        ) as mock_get_scorer,
+    ):
+        scorers = resolve_scorers(["Safety", "CustomScorer"], "experiment_123")
+
+        assert len(scorers) == 2
+        assert scorers[0] == mock_builtin
+        assert scorers[1] == mock_registered
+        # Verify mocks were called as expected
+        mock_get_all_scorers.assert_called_once()
+        mock_get_scorer.assert_called_once_with(name="CustomScorer", experiment_id="experiment_123")
+
+
+def test_resolve_scorer_not_found_raises_error():
+    with (
+        mock.patch(
+            "mlflow.cli.genai_eval_utils.get_all_scorers", return_value=[]
+        ) as mock_get_all_scorers,
+        mock.patch(
+            "mlflow.cli.genai_eval_utils.get_scorer",
+            side_effect=MlflowException("Not found"),
+        ) as mock_get_scorer,
+    ):
+        with pytest.raises(click.UsageError, match="Could not identify Scorer 'UnknownScorer'"):
+            resolve_scorers(["UnknownScorer"], "experiment_123")
+
+        # Verify mocks were called as expected
+        mock_get_all_scorers.assert_called_once()
+        mock_get_scorer.assert_called_once_with(
+            name="UnknownScorer", experiment_id="experiment_123"
+        )
+
+
+def test_resolve_empty_scorers_raises_error():
+    with pytest.raises(click.UsageError, match="No valid scorers"):
+        resolve_scorers([], "experiment_123")
+
+
+# Tests for extract_assessments_from_results function
+
+
+def test_extract_with_matching_run_id():
+    results_df = pd.DataFrame([
+        {
+            "trace_id": "tr-abc123",
+            "assessments": [
+                {
+                    "assessment_name": "RelevanceToQuery",
+                    "feedback": {"value": "yes"},
+                    "rationale": "The answer is relevant",
+                    "metadata": {AssessmentMetadataKey.SOURCE_RUN_ID: "run-123"},
+                }
+            ],
+        }
+    ])
+
+    result = extract_assessments_from_results(results_df, "run-123")
+
+    expected = [
+        EvalResult(
+            trace_id="tr-abc123",
+            assessments=[
+                Assessment(
+                    name="RelevanceToQuery",
+                    result="yes",
+                    rationale="The answer is relevant",
+                )
+            ],
+        )
+    ]
+    assert result == expected
+
+
+def test_extract_with_different_assessment_name():
+    results_df = pd.DataFrame([
+        {
+            "trace_id": "tr-abc123",
+            "assessments": [
+                {
+                    "assessment_name": "relevance_to_query",
+                    "feedback": {"value": "yes"},
+                    "rationale": "Relevant answer",
+                    "metadata": {AssessmentMetadataKey.SOURCE_RUN_ID: "run-123"},
+                }
+            ],
+        }
+    ])
+
+    result = extract_assessments_from_results(results_df, "run-123")
+
+    expected = [
+        EvalResult(
+            trace_id="tr-abc123",
+            assessments=[
+                Assessment(
+                    name="relevance_to_query",
+                    result="yes",
+                    rationale="Relevant answer",
+                )
+            ],
+        )
+    ]
+    assert result == expected
+
+
+def test_extract_filter_out_assessments_with_different_run_id():
+    results_df = pd.DataFrame([
+        {
+            "trace_id": "tr-abc123",
+            "assessments": [
+                {
+                    "assessment_name": "RelevanceToQuery",
+                    "feedback": {"value": "yes"},
+                    "rationale": "Current evaluation",
+                    "metadata": {AssessmentMetadataKey.SOURCE_RUN_ID: "run-123"},
+                },
+                {
+                    "assessment_name": "Safety",
+                    "feedback": {"value": "yes"},
+                    "rationale": "Old evaluation",
+                    "metadata": {AssessmentMetadataKey.SOURCE_RUN_ID: "run-456"},
+                },
+            ],
+        }
+    ])
+
+    result = extract_assessments_from_results(results_df, "run-123")
+
+    expected = [
+        EvalResult(
+            trace_id="tr-abc123",
+            assessments=[
+                Assessment(
+                    name="RelevanceToQuery",
+                    result="yes",
+                    rationale="Current evaluation",
+                )
+            ],
+        )
+    ]
+    assert result == expected
+
+
+def test_extract_no_assessments_for_run_id():
+    results_df = pd.DataFrame([
+        {
+            "trace_id": "tr-abc123",
+            "assessments": [
+                {
+                    "assessment_name": "RelevanceToQuery",
+                    "metadata": {AssessmentMetadataKey.SOURCE_RUN_ID: "run-456"},
+                }
+            ],
+        }
+    ])
+
+    result = extract_assessments_from_results(results_df, "run-123")
+
+    assert len(result) == 1
+    assert len(result[0].assessments) == 1
+    assert result[0].assessments[0].result is None
+    assert result[0].assessments[0].rationale is None
+    assert result[0].assessments[0].error is not None
+
+
+def test_extract_multiple_assessments_from_same_run():
+    results_df = pd.DataFrame([
+        {
+            "trace_id": "tr-abc123",
+            "assessments": [
+                {
+                    "assessment_name": "RelevanceToQuery",
+                    "feedback": {"value": "yes"},
+                    "rationale": "Relevant",
+                    "metadata": {AssessmentMetadataKey.SOURCE_RUN_ID: "run-123"},
+                },
+                {
+                    "assessment_name": "Safety",
+                    "feedback": {"value": "yes"},
+                    "rationale": "Safe",
+                    "metadata": {AssessmentMetadataKey.SOURCE_RUN_ID: "run-123"},
+                },
+            ],
+        }
+    ])
+
+    result = extract_assessments_from_results(results_df, "run-123")
+
+    expected = [
+        EvalResult(
+            trace_id="tr-abc123",
+            assessments=[
+                Assessment(
+                    name="RelevanceToQuery",
+                    result="yes",
+                    rationale="Relevant",
+                ),
+                Assessment(
+                    name="Safety",
+                    result="yes",
+                    rationale="Safe",
+                ),
+            ],
+        )
+    ]
+    assert result == expected
+
+
+def test_extract_no_assessments_on_trace_shows_error():
+    results_df = pd.DataFrame([{"trace_id": "tr-abc123", "assessments": []}])
+
+    result = extract_assessments_from_results(results_df, "run-123")
+
+    assert len(result) == 1
+    assert len(result[0].assessments) == 1
+    assert result[0].assessments[0].error == "No assessments found on trace"

tests/cli/test_scorers.py

@@ -0,0 +1,772 @@
+import json
+from typing import Any
+from unittest.mock import patch
+
+import pytest
+from click.testing import CliRunner
+
+import mlflow
+from mlflow.cli.scorers import commands
+from mlflow.exceptions import MlflowException
+from mlflow.genai.scorers import get_all_scorers, list_scorers, scorer
+from mlflow.utils.string_utils import _create_table
+
+
+@pytest.fixture
+def mock_databricks_environment():
+    with (
+        patch("mlflow.genai.scorers.base.is_databricks_uri", return_value=True),
+    ):
+        yield
+
+
+@pytest.fixture
+def runner():
+    return CliRunner(catch_exceptions=False)
+
+
+@pytest.fixture
+def experiment():
+    """Create a test experiment."""
+    experiment_id = mlflow.create_experiment(
+        f"test_scorers_cli_{mlflow.utils.time.get_current_time_millis()}"
+    )
+    yield experiment_id
+    mlflow.delete_experiment(experiment_id)
+
+
+@pytest.fixture
+def correctness_scorer():
+    """Create a correctness scorer."""
+
+    @scorer
+    def _correctness_scorer(outputs) -> bool:
+        return len(outputs) > 0
+
+    return _correctness_scorer
+
+
+@pytest.fixture
+def safety_scorer():
+    """Create a safety scorer."""
+
+    @scorer
+    def _safety_scorer(outputs) -> bool:
+        return len(outputs) > 0
+
+    return _safety_scorer
+
+
+@pytest.fixture
+def relevance_scorer():
+    """Create a relevance scorer."""
+
+    @scorer
+    def _relevance_scorer(outputs) -> bool:
+        return len(outputs) > 0
+
+    return _relevance_scorer
+
+
+@pytest.fixture
+def generic_scorer():
+    """Create a generic test scorer."""
+
+    @scorer
+    def _generic_scorer(outputs) -> bool:
+        return True
+
+    return _generic_scorer
+
+
+def test_commands_group_exists():
+    assert commands.name == "scorers"
+    assert commands.help is not None
+
+
+def test_list_command_params():
+    list_cmd = next((cmd for cmd in commands.commands.values() if cmd.name == "list"), None)
+    assert list_cmd is not None
+    param_names = {p.name for p in list_cmd.params}
+    assert param_names == {"experiment_id", "builtin", "output"}
+
+
+def test_list_scorers_table_output(
+    runner: CliRunner,
+    experiment: str,
+    correctness_scorer: Any,
+    safety_scorer: Any,
+    relevance_scorer: Any,
+    mock_databricks_environment: Any,
+):
+    correctness_scorer.register(experiment_id=experiment, name="Correctness")
+    safety_scorer.register(experiment_id=experiment, name="Safety")
+    relevance_scorer.register(experiment_id=experiment, name="RelevanceToQuery")
+
+    result = runner.invoke(commands, ["list", "--experiment-id", experiment])
+
+    assert result.exit_code == 0
+
+    # Construct expected table output (scorers are returned in alphabetical order)
+    # Note: click.echo() adds a trailing newline
+    expected_table = (
+        _create_table(
+            [["Correctness", ""], ["RelevanceToQuery", ""], ["Safety", ""]],
+            headers=["Scorer Name", "Description"],
+        )
+        + "\n"
+    )
+    assert result.output == expected_table
+
+
+def test_list_scorers_json_output(
+    runner: CliRunner,
+    experiment: str,
+    correctness_scorer: Any,
+    safety_scorer: Any,
+    relevance_scorer: Any,
+    mock_databricks_environment: Any,
+):
+    correctness_scorer.register(experiment_id=experiment, name="Correctness")
+    safety_scorer.register(experiment_id=experiment, name="Safety")
+    relevance_scorer.register(experiment_id=experiment, name="RelevanceToQuery")
+
+    result = runner.invoke(commands, ["list", "--experiment-id", experiment, "--output", "json"])
+
+    assert result.exit_code == 0
+    output_json = json.loads(result.output)
+    expected_scorers = [
+        {"name": "Correctness", "description": None},
+        {"name": "RelevanceToQuery", "description": None},
+        {"name": "Safety", "description": None},
+    ]
+    assert output_json["scorers"] == expected_scorers
+
+
+@pytest.mark.parametrize(
+    ("output_format", "expected_output"),
+    [
+        ("table", ""),
+        ("json", {"scorers": []}),
+    ],
+)
+def test_list_scorers_empty_experiment(
+    runner: CliRunner, experiment: str, output_format: str, expected_output: Any
+):
+    args = ["list", "--experiment-id", experiment]
+    if output_format == "json":
+        args.extend(["--output", "json"])
+
+    result = runner.invoke(commands, args)
+    assert result.exit_code == 0
+
+    if output_format == "json":
+        output_json = json.loads(result.output)
+        assert output_json == expected_output
+    else:
+        # Empty table produces minimal output
+        assert result.output.strip() == expected_output
+
+
+def test_list_scorers_with_experiment_id_env_var(
+    runner: CliRunner, experiment: str, correctness_scorer: Any, mock_databricks_environment: Any
+):
+    correctness_scorer.register(experiment_id=experiment, name="Correctness")
+
+    result = runner.invoke(commands, ["list"], env={"MLFLOW_EXPERIMENT_ID": experiment})
+
+    assert result.exit_code == 0
+    assert "Correctness" in result.output
+
+
+def test_list_scorers_missing_experiment_id(runner: CliRunner):
+    result = runner.invoke(commands, ["list"])
+
+    assert result.exit_code != 0
+    assert "experiment-id" in result.output.lower() or "experiment_id" in result.output.lower()
+
+
+def test_list_scorers_invalid_output_format(runner: CliRunner, experiment: str):
+    result = runner.invoke(commands, ["list", "--experiment-id", experiment, "--output", "invalid"])
+
+    assert result.exit_code != 0
+    assert "invalid" in result.output.lower() or "choice" in result.output.lower()
+
+
+def test_list_scorers_special_characters_in_names(
+    runner: CliRunner, experiment: str, generic_scorer: Any, mock_databricks_environment: Any
+):
+    generic_scorer.register(experiment_id=experiment, name="Scorer With Spaces")
+    generic_scorer.register(experiment_id=experiment, name="Scorer.With.Dots")
+    generic_scorer.register(experiment_id=experiment, name="Scorer-With-Dashes")
+    generic_scorer.register(experiment_id=experiment, name="Scorer_With_Underscores")
+
+    result = runner.invoke(commands, ["list", "--experiment-id", experiment])
+
+    assert result.exit_code == 0
+    assert "Scorer With Spaces" in result.output
+    assert "Scorer.With.Dots" in result.output
+    assert "Scorer-With-Dashes" in result.output
+    assert "Scorer_With_Underscores" in result.output
+
+
+@pytest.mark.parametrize(
+    "output_format",
+    ["table", "json"],
+)
+def test_list_scorers_single_scorer(
+    runner: CliRunner,
+    experiment: str,
+    generic_scorer: Any,
+    output_format: str,
+    mock_databricks_environment: Any,
+):
+    generic_scorer.register(experiment_id=experiment, name="OnlyScorer")
+
+    args = ["list", "--experiment-id", experiment]
+    if output_format == "json":
+        args.extend(["--output", "json"])
+
+    result = runner.invoke(commands, args)
+    assert result.exit_code == 0
+
+    if output_format == "json":
+        output_json = json.loads(result.output)
+        assert output_json == {"scorers": [{"name": "OnlyScorer", "description": None}]}
+    else:
+        assert "OnlyScorer" in result.output
+
+
+@pytest.mark.parametrize(
+    "output_format",
+    ["table", "json"],
+)
+def test_list_scorers_long_names(
+    runner: CliRunner,
+    experiment: str,
+    generic_scorer: Any,
+    output_format: str,
+    mock_databricks_environment: Any,
+):
+    long_name = "VeryLongScorerNameThatShouldNotBeTruncatedEvenIfItIsReallyReallyLong"
+    generic_scorer.register(experiment_id=experiment, name=long_name)
+
+    args = ["list", "--experiment-id", experiment]
+    if output_format == "json":
+        args.extend(["--output", "json"])
+
+    result = runner.invoke(commands, args)
+    assert result.exit_code == 0
+
+    if output_format == "json":
+        output_json = json.loads(result.output)
+        assert output_json == {"scorers": [{"name": long_name, "description": None}]}
+    else:
+        # Full name should be present
+        assert long_name in result.output
+
+
+def test_list_scorers_with_descriptions(runner: CliRunner, experiment: str):
+    from mlflow.genai.judges import make_judge
+
+    judge1 = make_judge(
+        name="quality_judge",
+        instructions="Evaluate {{ outputs }}",
+        description="Evaluates response quality",
+        feedback_value_type=str,
+    )
+    judge1.register(experiment_id=experiment)
+
+    judge2 = make_judge(
+        name="safety_judge",
+        instructions="Check {{ outputs }}",
+        description="Checks for safety issues",
+        feedback_value_type=str,
+    )
+    judge2.register(experiment_id=experiment)
+
+    judge3 = make_judge(
+        name="no_desc_judge",
+        instructions="Evaluate {{ outputs }}",
+        feedback_value_type=str,
+    )
+    judge3.register(experiment_id=experiment)
+
+    result_json = runner.invoke(
+        commands, ["list", "--experiment-id", experiment, "--output", "json"]
+    )
+    assert result_json.exit_code == 0
+    output_json = json.loads(result_json.output)
+
+    assert len(output_json["scorers"]) == 3
+    scorers_by_name = {s["name"]: s for s in output_json["scorers"]}
+
+    assert scorers_by_name["no_desc_judge"]["description"] is None
+    assert scorers_by_name["quality_judge"]["description"] == "Evaluates response quality"
+    assert scorers_by_name["safety_judge"]["description"] == "Checks for safety issues"
+
+    result_table = runner.invoke(commands, ["list", "--experiment-id", experiment])
+    assert result_table.exit_code == 0
+    assert "Evaluates response quality" in result_table.output
+    assert "Checks for safety issues" in result_table.output
+
+
+def test_create_judge_basic(runner: CliRunner, experiment: str):
+    result = runner.invoke(
+        commands,
+        [
+            "register-llm-judge",
+            "--name",
+            "test_judge",
+            "--instructions",
+            "Evaluate {{ outputs }}",
+            "--experiment-id",
+            experiment,
+        ],
+    )
+
+    assert result.exit_code == 0
+    assert "Successfully created and registered judge scorer 'test_judge'" in result.output
+    assert experiment in result.output
+
+    # Verify judge was registered
+    scorers = list_scorers(experiment_id=experiment)
+    scorer_names = [s.name for s in scorers]
+    assert "test_judge" in scorer_names
+
+
+def test_create_judge_with_model(runner: CliRunner, experiment: str):
+    result = runner.invoke(
+        commands,
+        [
+            "register-llm-judge",
+            "--name",
+            "custom_model_judge",
+            "--instructions",
+            "Check {{ inputs }} and {{ outputs }}",
+            "--model",
+            "openai:/gpt-4",
+            "--experiment-id",
+            experiment,
+        ],
+    )
+
+    assert result.exit_code == 0
+    assert "Successfully created and registered" in result.output
+
+    # Verify judge was registered with correct model
+    scorers = list_scorers(experiment_id=experiment)
+    scorer_names = [s.name for s in scorers]
+    assert "custom_model_judge" in scorer_names
+
+    # Get the judge and verify it uses the specified model
+    judge = next(s for s in scorers if s.name == "custom_model_judge")
+    assert judge.model == "openai:/gpt-4"
+
+
+def test_create_judge_short_options(runner: CliRunner, experiment: str):
+    result = runner.invoke(
+        commands,
+        [
+            "register-llm-judge",
+            "-n",
+            "short_options_judge",
+            "-i",
+            "Evaluate {{ outputs }}",
+            "-x",
+            experiment,
+        ],
+    )
+
+    assert result.exit_code == 0
+    assert "Successfully created and registered" in result.output
+
+    # Verify judge was registered
+    scorers = list_scorers(experiment_id=experiment)
+    scorer_names = [s.name for s in scorers]
+    assert "short_options_judge" in scorer_names
+
+
+def test_create_judge_with_env_var(runner: CliRunner, experiment: str):
+    result = runner.invoke(
+        commands,
+        [
+            "register-llm-judge",
+            "--name",
+            "env_var_judge",
+            "--instructions",
+            "Check {{ outputs }}",
+        ],
+        env={"MLFLOW_EXPERIMENT_ID": experiment},
+    )
+
+    assert result.exit_code == 0
+    assert "Successfully created and registered" in result.output
+
+    # Verify judge was registered
+    scorers = list_scorers(experiment_id=experiment)
+    scorer_names = [s.name for s in scorers]
+    assert "env_var_judge" in scorer_names
+
+
+@pytest.mark.parametrize(
+    ("args", "missing_param"),
+    [
+        (["--instructions", "test", "--experiment-id", "123"], "name"),
+        (["--name", "test", "--experiment-id", "123"], "instructions"),
+        (["--name", "test", "--instructions", "test"], "experiment-id"),
+    ],
+)
+def test_create_judge_missing_required_params(
+    runner: CliRunner, args: list[str], missing_param: str
+):
+    result = runner.invoke(commands, ["register-llm-judge"] + args)
+
+    assert result.exit_code != 0
+    # Click typically shows "Missing option" for required parameters
+    assert "missing" in result.output.lower() or "required" in result.output.lower()
+
+
+def test_create_judge_invalid_prompt(runner: CliRunner, experiment: str):
+    # Should raise MlflowException because make_judge validates that instructions
+    # contain at least one variable
+    with pytest.raises(MlflowException, match="[Tt]emplate.*variable"):
+        runner.invoke(
+            commands,
+            [
+                "register-llm-judge",
+                "--name",
+                "invalid_judge",
+                "--instructions",
+                "This has no template variables",
+                "--experiment-id",
+                experiment,
+            ],
+        )
+
+
+def test_create_judge_special_characters_in_name(runner: CliRunner, experiment: str):
+    # Verify experiment has no judges initially
+    scorers = list_scorers(experiment_id=experiment)
+    assert len(scorers) == 0
+
+    result = runner.invoke(
+        commands,
+        [
+            "register-llm-judge",
+            "--name",
+            "judge-with_special.chars",
+            "--instructions",
+            "Evaluate {{ outputs }}",
+            "--experiment-id",
+            experiment,
+        ],
+    )
+
+    assert result.exit_code == 0
+    assert "Successfully created and registered" in result.output
+
+    # Verify experiment has exactly one judge
+    scorers = list_scorers(experiment_id=experiment)
+    assert len(scorers) == 1
+    assert scorers[0].name == "judge-with_special.chars"
+
+
+def test_create_judge_duplicate_registration(runner: CliRunner, experiment: str):
+    # Create a judge
+    result1 = runner.invoke(
+        commands,
+        [
+            "register-llm-judge",
+            "--name",
+            "duplicate_judge",
+            "--instructions",
+            "Evaluate {{ outputs }}",
+            "--experiment-id",
+            experiment,
+        ],
+    )
+    assert result1.exit_code == 0
+
+    scorers = list_scorers(experiment_id=experiment)
+    assert len(scorers) == 1
+    assert scorers[0].name == "duplicate_judge"
+
+    # Register the same judge again with same name - should succeed (replaces the old one)
+    result2 = runner.invoke(
+        commands,
+        [
+            "register-llm-judge",
+            "--name",
+            "duplicate_judge",
+            "--instructions",
+            "Evaluate {{ outputs }}",
+            "--experiment-id",
+            experiment,
+        ],
+    )
+    assert result2.exit_code == 0
+
+    # Verify there is still only one judge (the new one replaced the old one)
+    scorers = list_scorers(experiment_id=experiment)
+    assert len(scorers) == 1
+    assert scorers[0].name == "duplicate_judge"
+
+
+def test_create_judge_with_description(runner: CliRunner, experiment: str):
+    description = "Evaluates response quality and relevance"
+    result = runner.invoke(
+        commands,
+        [
+            "register-llm-judge",
+            "--name",
+            "judge_with_desc",
+            "--instructions",
+            "Evaluate {{ outputs }}",
+            "--description",
+            description,
+            "--experiment-id",
+            experiment,
+        ],
+    )
+
+    assert result.exit_code == 0
+    assert "Successfully created and registered" in result.output
+
+    scorers = list_scorers(experiment_id=experiment)
+    assert len(scorers) == 1
+    judge = scorers[0]
+    assert judge.name == "judge_with_desc"
+    assert judge.description == description
+
+
+def test_create_judge_with_description_short_flag(runner: CliRunner, experiment: str):
+    description = "Checks for PII in outputs"
+    result = runner.invoke(
+        commands,
+        [
+            "register-llm-judge",
+            "-n",
+            "pii_judge",
+            "-i",
+            "Check {{ outputs }}",
+            "-d",
+            description,
+            "-x",
+            experiment,
+        ],
+    )
+
+    assert result.exit_code == 0
+
+    scorers = list_scorers(experiment_id=experiment)
+    judge = next(s for s in scorers if s.name == "pii_judge")
+    assert judge.description == description
+
+
+@pytest.mark.parametrize("output_format", ["table", "json"])
+def test_list_builtin_scorers_output_formats(runner, output_format):
+    args = ["list", "--builtin"]
+    if output_format == "json":
+        args.extend(["--output", "json"])
+
+    result = runner.invoke(commands, args)
+    assert result.exit_code == 0
+
+    if output_format == "json":
+        data = json.loads(result.output)
+        assert "scorers" in data
+        assert isinstance(data["scorers"], list)
+        assert len(data["scorers"]) > 0
+
+        # Verify each scorer has required fields
+        for scorer_item in data["scorers"]:
+            assert "name" in scorer_item
+            assert "description" in scorer_item
+
+        # Verify some builtin scorer names appear
+        scorer_names = [s["name"] for s in data["scorers"]]
+        assert "correctness" in scorer_names
+        assert "relevance_to_query" in scorer_names
+        assert "completeness" in scorer_names
+    else:
+        # Verify table headers
+        assert "Scorer Name" in result.output
+        assert "Description" in result.output
+
+        # Verify some builtin scorer names appear
+        assert "correctness" in result.output
+        assert "relevance_to_query" in result.output
+        assert "completeness" in result.output
+
+
+def test_list_builtin_scorers_short_flag(runner):
+    result = runner.invoke(commands, ["list", "-b"])
+    assert result.exit_code == 0
+    assert "Scorer Name" in result.output
+
+
+def test_list_builtin_scorers_shows_all_available_scorers(runner):
+    result = runner.invoke(commands, ["list", "--builtin", "--output", "json"])
+    assert result.exit_code == 0
+
+    expected_scorers = get_all_scorers()
+    expected_names = {scorer.name for scorer in expected_scorers}
+
+    data = json.loads(result.output)
+    actual_names = {s["name"] for s in data["scorers"]}
+
+    assert actual_names == expected_names
+
+
+def test_list_scorers_mutually_exclusive_flags(runner, experiment):
+    result = runner.invoke(commands, ["list", "--builtin", "--experiment-id", experiment])
+    assert result.exit_code != 0
+    assert "Cannot specify both --builtin and --experiment-id" in result.output
+
+
+def test_list_scorers_requires_one_flag(runner):
+    result = runner.invoke(commands, ["list"])
+    assert result.exit_code != 0
+    assert "Must specify either --builtin or --experiment-id" in result.output
+
+
+def test_list_scorers_env_var_still_works(runner, experiment, monkeypatch):
+    monkeypatch.setenv("MLFLOW_EXPERIMENT_ID", experiment)
+    result = runner.invoke(commands, ["list"])
+    assert result.exit_code == 0
+
+
+def test_create_judge_with_base_url(runner: CliRunner, experiment: str):
+    result = runner.invoke(
+        commands,
+        [
+            "register-llm-judge",
+            "--name",
+            "proxy_judge",
+            "--instructions",
+            "Evaluate {{ outputs }}",
+            "--model",
+            "openai:/gpt-4",
+            "--base-url",
+            "http://my-proxy:8080/v1",
+            "--experiment-id",
+            experiment,
+        ],
+    )
+
+    assert result.exit_code == 0
+    assert "Successfully created and registered" in result.output
+
+    # base_url is not persisted, so the registered judge won't have it
+    scorers = list_scorers(experiment_id=experiment)
+    assert any(s.name == "proxy_judge" for s in scorers)
+
+
+def test_create_judge_with_extra_headers(runner: CliRunner, experiment: str):
+    result = runner.invoke(
+        commands,
+        [
+            "register-llm-judge",
+            "--name",
+            "headers_judge",
+            "--instructions",
+            "Evaluate {{ outputs }}",
+            "--model",
+            "openai:/gpt-4",
+            "--extra-headers",
+            '{"X-Api-Key": "secret", "X-Org": "my-org"}',
+            "--experiment-id",
+            experiment,
+        ],
+    )
+
+    assert result.exit_code == 0
+    assert "Successfully created and registered" in result.output
+
+    scorers = list_scorers(experiment_id=experiment)
+    assert any(s.name == "headers_judge" for s in scorers)
+
+
+def test_create_judge_with_base_url_and_extra_headers(runner: CliRunner, experiment: str):
+    result = runner.invoke(
+        commands,
+        [
+            "register-llm-judge",
+            "--name",
+            "full_judge",
+            "--instructions",
+            "Evaluate {{ outputs }}",
+            "--model",
+            "openai:/gpt-4",
+            "--base-url",
+            "http://proxy:9090",
+            "--extra-headers",
+            '{"Authorization": "Bearer token"}',
+            "--experiment-id",
+            experiment,
+        ],
+    )
+
+    assert result.exit_code == 0
+    assert "Successfully created and registered" in result.output
+
+
+def test_create_judge_invalid_extra_headers_json(runner: CliRunner, experiment: str):
+    result = runner.invoke(
+        commands,
+        [
+            "register-llm-judge",
+            "--name",
+            "bad_json_judge",
+            "--instructions",
+            "Evaluate {{ outputs }}",
+            "--extra-headers",
+            "not valid json",
+            "--experiment-id",
+            experiment,
+        ],
+    )
+
+    assert result.exit_code != 0
+    assert "Invalid JSON" in result.output
+
+
+def test_create_judge_extra_headers_not_dict(runner: CliRunner, experiment: str):
+    result = runner.invoke(
+        commands,
+        [
+            "register-llm-judge",
+            "--name",
+            "array_headers_judge",
+            "--instructions",
+            "Evaluate {{ outputs }}",
+            "--extra-headers",
+            '["not", "a", "dict"]',
+            "--experiment-id",
+            experiment,
+        ],
+    )
+
+    assert result.exit_code != 0
+    assert "Expected a JSON object" in result.output
+
+
+def test_create_judge_extra_headers_non_string_values(runner: CliRunner, experiment: str):
+    result = runner.invoke(
+        commands,
+        [
+            "register-llm-judge",
+            "--name",
+            "non_string_headers_judge",
+            "--instructions",
+            "Evaluate {{ outputs }}",
+            "--extra-headers",
+            '{"Authorization": 123}',
+            "--experiment-id",
+            experiment,
+        ],
+    )
+
+    assert result.exit_code != 0
+    assert "must all be strings" in result.output

tests/cli/test_skills.py

@@ -0,0 +1,79 @@
+from pathlib import Path
+from unittest import mock
+
+import pytest
+from click.testing import CliRunner
+
+from mlflow.assistant.skill_installer import BundledSkill
+from mlflow.cli.skills import commands
+
+
+@pytest.fixture
+def runner():
+    return CliRunner()
+
+
+@pytest.fixture
+def mock_bundled_skills():
+    with mock.patch("mlflow.cli.skills.list_bundled_skills", return_value=[]) as m:
+        yield m
+
+
+def test_list_command_renders_skills(runner, mock_bundled_skills):
+    skills = [
+        BundledSkill(name="alpha-skill", description="Does alpha things.", path=Path("/s/alpha")),
+        BundledSkill(name="beta-skill", description="", path=Path("/s/beta")),
+    ]
+    mock_bundled_skills.return_value = skills
+
+    result = runner.invoke(commands, ["list"])
+
+    assert result.exit_code == 0
+    assert "alpha-skill" in result.output
+    assert "Does alpha things." in result.output
+    assert "beta-skill" in result.output
+    assert str(Path("/s/alpha")) in result.output
+    assert str(Path("/s/beta")) in result.output
+
+
+def test_list_command_handles_no_skills(runner, mock_bundled_skills):
+    result = runner.invoke(commands, ["list"])
+
+    assert result.exit_code == 0
+    assert "No MLflow skills found" in result.output
+    assert "git submodule update --init" in result.output
+
+
+def test_view_command_renders_skill(runner, mock_bundled_skills):
+    skills = [
+        BundledSkill(name="alpha-skill", description="Does alpha things.", path=Path("/s/alpha")),
+        BundledSkill(name="beta-skill", description="Does beta things.", path=Path("/s/beta")),
+    ]
+    mock_bundled_skills.return_value = skills
+
+    result = runner.invoke(commands, ["view", "alpha-skill"])
+
+    assert result.exit_code == 0
+    assert "alpha-skill" in result.output
+    assert "Does alpha things." in result.output
+    assert str(Path("/s/alpha")) in result.output
+    assert "beta-skill" not in result.output
+
+
+def test_view_command_skill_not_found(runner, mock_bundled_skills):
+    result = runner.invoke(commands, ["view", "missing-skill"])
+
+    assert result.exit_code != 0
+    assert "Skill missing-skill not found." in result.output
+
+
+def test_view_command_skill_without_description(runner, mock_bundled_skills):
+    skills = [BundledSkill(name="alpha-skill", description="", path=Path("/s/alpha"))]
+    mock_bundled_skills.return_value = skills
+
+    result = runner.invoke(commands, ["view", "alpha-skill"])
+
+    assert result.exit_code == 0
+    assert "alpha-skill" in result.output
+    assert str(Path("/s/alpha")) in result.output
+    assert "\n  \n" not in result.output

tests/cli/test_traces.py

@@ -0,0 +1,228 @@
+import json
+import logging
+from unittest import mock
+
+import pytest
+from click.testing import CliRunner
+
+from mlflow.cli.traces import commands
+from mlflow.entities import (
+    AssessmentSourceType,
+    MlflowExperimentLocation,
+    Trace,
+    TraceData,
+    TraceInfo,
+    TraceLocation,
+    TraceLocationType,
+    TraceState,
+)
+from mlflow.store.entities.paged_list import PagedList
+
+
+@pytest.fixture(autouse=True)
+def suppress_logging():
+    """Suppress logging for all tests."""
+    # Suppress logging
+    original_root = logging.root.level
+    original_mlflow = logging.getLogger("mlflow").level
+    original_alembic = logging.getLogger("alembic").level
+
+    logging.root.setLevel(logging.CRITICAL)
+    logging.getLogger("mlflow").setLevel(logging.CRITICAL)
+    logging.getLogger("alembic").setLevel(logging.CRITICAL)
+
+    yield
+
+    # Restore original logging levels
+    logging.root.setLevel(original_root)
+    logging.getLogger("mlflow").setLevel(original_mlflow)
+    logging.getLogger("alembic").setLevel(original_alembic)
+
+
+@pytest.fixture
+def runner():
+    """Provide a CLI runner for testing."""
+    return CliRunner(catch_exceptions=False)
+
+
+def test_commands_group_exists():
+    assert commands.name == "traces"
+    assert commands.help is not None
+
+
+def test_search_command_params():
+    search_cmd = next((cmd for cmd in commands.commands.values() if cmd.name == "search"), None)
+    assert search_cmd is not None
+    param_names = [p.name for p in search_cmd.params]
+    assert "experiment_id" in param_names
+    assert "filter_string" in param_names
+    assert "max_results" in param_names
+    assert "order_by" in param_names
+    assert "page_token" in param_names
+    assert "output" in param_names
+    assert "extract_fields" in param_names
+
+
+def test_get_command_params():
+    get_cmd = next((cmd for cmd in commands.commands.values() if cmd.name == "get"), None)
+    assert get_cmd is not None
+    param_names = [p.name for p in get_cmd.params]
+    assert "trace_id" in param_names
+    assert "extract_fields" in param_names
+
+
+def test_assessment_source_type_choices():
+    log_feedback_cmd = next(
+        (cmd for cmd in commands.commands.values() if cmd.name == "log-feedback"), None
+    )
+    assert log_feedback_cmd is not None
+
+    source_type_param = next(
+        (param for param in log_feedback_cmd.params if param.name == "source_type"), None
+    )
+    assert source_type_param is not None
+    assert AssessmentSourceType.HUMAN in source_type_param.type.choices
+    assert AssessmentSourceType.LLM_JUDGE in source_type_param.type.choices
+    assert AssessmentSourceType.CODE in source_type_param.type.choices
+
+
+def test_search_command_with_fields(runner):
+    trace_location = TraceLocation(
+        type=TraceLocationType.MLFLOW_EXPERIMENT,
+        mlflow_experiment=MlflowExperimentLocation(experiment_id="1"),
+    )
+    trace = Trace(
+        info=TraceInfo(
+            trace_id="tr-123",
+            state=TraceState.OK,
+            request_time=1700000000000,
+            execution_duration=1234,
+            request_preview="test request",
+            response_preview="test response",
+            trace_location=trace_location,
+        ),
+        data=TraceData(spans=[]),
+    )
+
+    mock_result = PagedList([trace], None)
+
+    with mock.patch("mlflow.cli.traces.TracingClient") as mock_client:
+        mock_client.return_value.search_traces.return_value = mock_result
+        result = runner.invoke(
+            commands,
+            ["search", "--experiment-id", "1", "--extract-fields", "info.trace_id,info.state"],
+        )
+
+        assert result.exit_code == 0
+        assert "tr-123" in result.output
+        assert "OK" in result.output
+
+
+def test_get_command_with_fields(runner):
+    trace_location = TraceLocation(
+        type=TraceLocationType.MLFLOW_EXPERIMENT,
+        mlflow_experiment=MlflowExperimentLocation(experiment_id="1"),
+    )
+    trace = Trace(
+        info=TraceInfo(
+            trace_id="tr-123",
+            state=TraceState.OK,
+            trace_location=trace_location,
+            request_time=1700000000000,
+            execution_duration=1234,
+        ),
+        data=TraceData(spans=[]),
+    )
+
+    with mock.patch("mlflow.cli.traces.TracingClient") as mock_client:
+        mock_client.return_value.get_trace.return_value = trace
+        result = runner.invoke(
+            commands,
+            ["get", "--trace-id", "tr-123", "--extract-fields", "info.trace_id"],
+        )
+
+        assert result.exit_code == 0
+        output_json = json.loads(result.output)
+        assert output_json == {"info": {"trace_id": "tr-123"}}
+
+
+def test_delete_command(runner):
+    with mock.patch("mlflow.cli.traces.TracingClient") as mock_client:
+        mock_client.return_value.delete_traces.return_value = 5
+        result = runner.invoke(
+            commands,
+            ["delete", "--experiment-id", "1", "--trace-ids", "tr-1,tr-2,tr-3"],
+        )
+
+        assert result.exit_code == 0
+        assert "Deleted 5 trace(s)" in result.output
+
+
+def test_field_validation_error(runner):
+    trace_location = TraceLocation(
+        type=TraceLocationType.MLFLOW_EXPERIMENT,
+        mlflow_experiment=MlflowExperimentLocation(experiment_id="1"),
+    )
+    trace = Trace(
+        info=TraceInfo(
+            trace_id="tr-123",
+            trace_location=trace_location,
+            request_time=1700000000000,
+            execution_duration=1234,
+            state=TraceState.OK,
+        ),
+        data=TraceData(spans=[]),
+    )
+
+    mock_result = PagedList([trace], None)
+
+    with mock.patch("mlflow.cli.traces.TracingClient") as mock_client:
+        mock_client.return_value.search_traces.return_value = mock_result
+        result = runner.invoke(
+            commands,
+            ["search", "--experiment-id", "1", "--extract-fields", "invalid.field"],
+        )
+
+        assert result.exit_code != 0
+        assert "Invalid field path" in result.output
+        assert "--verbose" in result.output
+
+
+def test_field_validation_error_verbose_mode(runner):
+    trace_location = TraceLocation(
+        type=TraceLocationType.MLFLOW_EXPERIMENT,
+        mlflow_experiment=MlflowExperimentLocation(experiment_id="1"),
+    )
+    trace = Trace(
+        info=TraceInfo(
+            trace_id="tr-123",
+            state=TraceState.OK,
+            request_time=1700000000000,
+            trace_location=trace_location,
+            execution_duration=1234,
+        ),
+        data=TraceData(spans=[]),
+    )
+
+    mock_result = PagedList([trace], None)
+
+    with mock.patch("mlflow.cli.traces.TracingClient") as mock_client:
+        mock_client.return_value.search_traces.return_value = mock_result
+        result = runner.invoke(
+            commands,
+            [
+                "search",
+                "--experiment-id",
+                "1",
+                "--extract-fields",
+                "invalid.field",
+                "--verbose",
+            ],
+        )
+
+        assert result.exit_code != 0
+        assert "Invalid field path" in result.output
+        assert "info.trace_id" in result.output
+        assert "info.state" in result.output
+        assert "info.request_time" in result.output
+        assert "Tip: Use --verbose" not in result.output