mlflow--mlflow/tests/test_cli.py

import os
import shutil
import subprocess
import sys
import tempfile
import time
from pathlib import Path
from unittest import mock
from urllib.parse import unquote, urlparse
from urllib.request import url2pathname

import click
import numpy as np
import pytest
import requests
from botocore.stub import Stubber
from click.testing import CliRunner

import mlflow
from mlflow.cli import cli, doctor, gc, server
from mlflow.data import numpy_dataset
from mlflow.entities import Metric, ViewType
from mlflow.entities.logged_model import LoggedModelParameter, LoggedModelTag
from mlflow.environment_variables import (
    MLFLOW_ENABLE_WORKSPACES,
    MLFLOW_TRACE_ARCHIVAL_CONFIG,
    MLFLOW_WORKSPACE_STORE_URI,
)
from mlflow.exceptions import MlflowException
from mlflow.server import handlers
from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
from mlflow.store.jobs.sqlalchemy_store import SqlAlchemyJobStore
from mlflow.store.tracking.dbmodels.models import (
    SqlLoggedModelMetric,
    SqlLoggedModelParam,
    SqlLoggedModelTag,
)
from mlflow.store.tracking.file_store import FileStore
from mlflow.store.tracking.sqlalchemy_store import SqlAlchemyStore
from mlflow.utils.os import is_windows
from mlflow.utils.rest_utils import augmented_raise_for_status
from mlflow.utils.time import get_current_time_millis

from tests.helper_functions import (
    get_safe_port,
    kill_process_tree,
)
from tests.tracking.integration_test_utils import _await_server_up_or_die


@pytest.mark.parametrize("command", ["server"])
def test_mlflow_server_command(command):
    port = get_safe_port()
    with subprocess.Popen([
        sys.executable,
        "-m",
        "mlflow",
        command,
        "--port",
        str(port),
    ]) as process:
        try:
            _await_server_up_or_die(port)
            resp = requests.get(f"http://localhost:{port}/health")
            augmented_raise_for_status(resp)
            assert resp.text == "OK"
        finally:
            kill_process_tree(process.pid)


def test_server_static_prefix_validation():
    with mock.patch("mlflow.server._run_server") as run_server_mock:
        CliRunner().invoke(server)
        run_server_mock.assert_called_once()
    with mock.patch("mlflow.server._run_server") as run_server_mock:
        CliRunner().invoke(server, ["--static-prefix", "/mlflow"])
        run_server_mock.assert_called_once()
    with mock.patch("mlflow.server._run_server") as run_server_mock:
        result = CliRunner().invoke(server, ["--static-prefix", "mlflow/"])
        assert "--static-prefix must begin with a '/'." in result.output
        run_server_mock.assert_not_called()
    with mock.patch("mlflow.server._run_server") as run_server_mock:
        result = CliRunner().invoke(server, ["--static-prefix", "/mlflow/"])
        assert "--static-prefix should not end with a '/'." in result.output
        run_server_mock.assert_not_called()


def test_server_cli_fails_when_workspace_env_set():
    env = os.environ.copy()
    # Trigger server-mode guard in mlflow.server.__init__
    env["_MLFLOW_SERVER_SERVE_ARTIFACTS"] = "1"
    env["MLFLOW_WORKSPACE"] = "team-a"

    result = subprocess.run(
        [sys.executable, "-m", "mlflow", "server", "--port", "0"],
        capture_output=True,
        text=True,
        env=env,
    )
    assert result.returncode != 0
    assert "MLFLOW_WORKSPACE=team-a is client-only" in result.stderr


def test_server_uvicorn_options():
    with mock.patch("mlflow.server._run_server") as run_server_mock:
        # Test default behavior (uvicorn should be used when no server options specified)
        CliRunner().invoke(server)
        run_server_mock.assert_called_once_with(
            file_store_path=mock.ANY,
            read_replica_backend_store_uri=mock.ANY,
            registry_store_uri=mock.ANY,
            default_artifact_root=mock.ANY,
            serve_artifacts=mock.ANY,
            artifacts_only=mock.ANY,
            artifacts_destination=mock.ANY,
            host="127.0.0.1",
            port=5000,
            static_prefix=None,
            workers=None,
            gunicorn_opts=None,
            waitress_opts=None,
            expose_prometheus=None,
            app_name=None,
            uvicorn_opts=None,
            env_file=None,
            secrets_cache_ttl=60,
            secrets_cache_max_size=1000,
        )

    with mock.patch("mlflow.server._run_server") as run_server_mock:
        # Test with uvicorn-opts - use different options than dev mode
        CliRunner().invoke(server, ["--uvicorn-opts", "--loop asyncio --limit-concurrency 100"])
        run_server_mock.assert_called_once_with(
            file_store_path=mock.ANY,
            read_replica_backend_store_uri=mock.ANY,
            registry_store_uri=mock.ANY,
            default_artifact_root=mock.ANY,
            serve_artifacts=mock.ANY,
            artifacts_only=mock.ANY,
            artifacts_destination=mock.ANY,
            host="127.0.0.1",
            port=5000,
            static_prefix=None,
            workers=None,
            gunicorn_opts=None,
            waitress_opts=None,
            expose_prometheus=None,
            app_name=None,
            uvicorn_opts="--loop asyncio --limit-concurrency 100",
            env_file=None,
            secrets_cache_ttl=60,
            secrets_cache_max_size=1000,
        )


@pytest.mark.skipif(is_windows(), reason="--dev mode is not supported on Windows")
def test_server_dev_mode():
    with mock.patch("mlflow.server._run_server") as run_server_mock:
        # Test with --dev flag (should set uvicorn opts)
        CliRunner().invoke(server, ["--dev"])
        run_server_mock.assert_called_once_with(
            file_store_path=mock.ANY,
            read_replica_backend_store_uri=mock.ANY,
            registry_store_uri=mock.ANY,
            default_artifact_root=mock.ANY,
            serve_artifacts=mock.ANY,
            artifacts_only=mock.ANY,
            artifacts_destination=mock.ANY,
            host="127.0.0.1",
            port=5000,
            static_prefix=None,
            workers=None,
            gunicorn_opts=None,
            waitress_opts=None,
            expose_prometheus=None,
            app_name=None,
            uvicorn_opts="--reload --log-level debug",
            env_file=None,
            secrets_cache_ttl=60,
            secrets_cache_max_size=1000,
        )


@pytest.mark.skipif(is_windows(), reason="Gunicorn is not supported on Windows")
def test_server_gunicorn_options():
    with mock.patch("mlflow.server._run_server") as run_server_mock:
        # Test that gunicorn-opts disables uvicorn
        CliRunner().invoke(server, ["--gunicorn-opts", "--timeout 120 --max-requests 1000"])
        run_server_mock.assert_called_once_with(
            file_store_path=mock.ANY,
            read_replica_backend_store_uri=mock.ANY,
            registry_store_uri=mock.ANY,
            default_artifact_root=mock.ANY,
            serve_artifacts=mock.ANY,
            artifacts_only=mock.ANY,
            artifacts_destination=mock.ANY,
            host="127.0.0.1",
            port=5000,
            static_prefix=None,
            workers=None,
            gunicorn_opts="--timeout 120 --max-requests 1000",
            waitress_opts=None,
            expose_prometheus=None,
            app_name=None,
            uvicorn_opts=None,
            env_file=None,
            secrets_cache_ttl=60,
            secrets_cache_max_size=1000,
        )

    # Test conflicting options
    result = CliRunner().invoke(
        server, ["--uvicorn-opts", "--reload", "--gunicorn-opts", "--log-level debug"]
    )
    assert result.exit_code != 0
    assert "Cannot specify multiple server options" in result.output


def test_server_initializes_backend_store_when_tracking_enabled():
    handlers._tracking_store = None
    handlers._model_registry_store = None
    runner = CliRunner()
    with runner.isolated_filesystem():
        with (
            mock.patch("mlflow.server.handlers.initialize_backend_stores") as init_backend_mock,
            mock.patch("mlflow.server._run_server") as run_server_mock,
        ):
            result = runner.invoke(server)
    assert result.exit_code == 0
    init_backend_mock.assert_called_once_with(
        mock.ANY, mock.ANY, mock.ANY, workspace_store_uri=None, read_replica_backend_store_uri=None
    )
    run_server_mock.assert_called_once()


def test_server_skips_backend_store_init_in_artifacts_only_mode():
    handlers._tracking_store = None
    handlers._model_registry_store = None
    runner = CliRunner()
    with runner.isolated_filesystem():
        with (
            mock.patch("mlflow.server.handlers.initialize_backend_stores") as init_backend_mock,
            mock.patch("mlflow.server._run_server") as run_server_mock,
        ):
            result = runner.invoke(server, ["--artifacts-only"])
    assert result.exit_code == 0
    init_backend_mock.assert_not_called()
    run_server_mock.assert_called_once()


def test_server_mlflow_artifacts_options():
    handlers._tracking_store = None
    handlers._model_registry_store = None
    with mock.patch(
        "mlflow.tracking._tracking_service.utils._has_existing_mlruns_data", return_value=False
    ):
        runner = CliRunner()
        with runner.isolated_filesystem():
            with mock.patch("mlflow.server._run_server") as run_server_mock:
                runner.invoke(server, ["--artifacts-only"])
                run_server_mock.assert_called_once()
            with mock.patch("mlflow.server._run_server") as run_server_mock:
                runner.invoke(server, ["--serve-artifacts"])
                run_server_mock.assert_called_once()
            with mock.patch("mlflow.server._run_server") as run_server_mock:
                runner.invoke(server, ["--no-serve-artifacts"])
                run_server_mock.assert_called_once()
            with mock.patch("mlflow.server._run_server") as run_server_mock:
                runner.invoke(server, ["--artifacts-only"])
                run_server_mock.assert_called_once()


def test_server_artifacts_only_conflicts_with_enable_workspaces():
    with pytest.raises(
        click.UsageError, match="--enable-workspaces cannot be combined with --artifacts-only"
    ):
        CliRunner().invoke(
            server,
            ["--artifacts-only", "--enable-workspaces"],
            catch_exceptions=False,
            standalone_mode=False,
        )


def _write_trace_archival_config(
    tmp_path: Path,
    *,
    enabled: bool = True,
    location: str = "s3://archive/default",
    retention: str = "30d",
    long_retention_allowlist: list[str] | None = None,
    interval_seconds: int | None = None,
    max_traces_per_pass: int | None = None,
) -> Path:
    lines = [
        "trace_archival:",
        f"  enabled: {'true' if enabled else 'false'}",
        f"  location: {location}",
        f"  retention: {retention}",
    ]
    if long_retention_allowlist is not None:
        lines.append("  long_retention_allowlist:")
        lines.extend(f'    - "{entry}"' for entry in long_retention_allowlist)
    if interval_seconds is not None:
        lines.append(f"  interval_seconds: {interval_seconds}")
    if max_traces_per_pass is not None:
        lines.append(f"  max_traces_per_pass: {max_traces_per_pass}")

    config_path = tmp_path / "trace-archival.yaml"
    config_path.parent.mkdir(parents=True, exist_ok=True)
    config_path.write_text("\n".join(lines) + "\n", encoding="utf-8")
    return config_path


def test_server_artifacts_only_conflicts_with_trace_archival_config(tmp_path):
    config_path = _write_trace_archival_config(tmp_path)

    with pytest.raises(
        click.UsageError,
        match="--trace-archival-config cannot be combined with --artifacts-only",
    ):
        CliRunner().invoke(
            server,
            ["--artifacts-only", "--trace-archival-config", str(config_path)],
            catch_exceptions=False,
            standalone_mode=False,
        )


def test_server_artifacts_only_conflicts_with_trace_archival_config_env_var(tmp_path):
    config_path = _write_trace_archival_config(tmp_path)
    MLFLOW_TRACE_ARCHIVAL_CONFIG.set(str(config_path))
    try:
        with pytest.raises(
            click.UsageError,
            match="--trace-archival-config cannot be combined with --artifacts-only",
        ):
            CliRunner().invoke(
                server,
                ["--artifacts-only"],
                catch_exceptions=False,
                standalone_mode=False,
            )
    finally:
        MLFLOW_TRACE_ARCHIVAL_CONFIG.unset()


def test_server_workspace_uri_sets_env_when_workspaces_enabled(tmp_path):
    handlers._tracking_store = None
    handlers._model_registry_store = None
    workspace_uri = f"sqlite:///{tmp_path / 'workspace.db'}"
    backend_uri = f"sqlite:///{tmp_path / 'backend.db'}"
    artifact_root_path = tmp_path / "artifacts"
    artifact_root_path.mkdir()
    artifact_root = artifact_root_path.as_uri()

    MLFLOW_WORKSPACE_STORE_URI.unset()
    MLFLOW_ENABLE_WORKSPACES.unset()
    MLFLOW_TRACE_ARCHIVAL_CONFIG.unset()

    try:
        with (
            mock.patch("mlflow.server._run_server") as run_server_mock,
            mock.patch("mlflow.server.handlers.initialize_backend_stores") as init_backend,
        ):
            result = CliRunner().invoke(
                server,
                [
                    "--enable-workspaces",
                    "--workspace-store-uri",
                    workspace_uri,
                    "--backend-store-uri",
                    backend_uri,
                    "--registry-store-uri",
                    backend_uri,
                    "--default-artifact-root",
                    artifact_root,
                ],
                catch_exceptions=False,
                standalone_mode=False,
            )
        assert result.exit_code == 0
        run_server_mock.assert_called_once()
        init_backend.assert_called_once_with(
            backend_uri,
            backend_uri,
            artifact_root,
            workspace_store_uri=workspace_uri,
            read_replica_backend_store_uri=None,
        )
        assert MLFLOW_WORKSPACE_STORE_URI.get() == workspace_uri
        assert MLFLOW_ENABLE_WORKSPACES.get() is True
    finally:
        MLFLOW_WORKSPACE_STORE_URI.unset()
        MLFLOW_ENABLE_WORKSPACES.unset()
        MLFLOW_TRACE_ARCHIVAL_CONFIG.unset()


def test_server_trace_archival_config_sets_env(tmp_path):
    handlers._tracking_store = None
    handlers._model_registry_store = None
    backend_uri = f"sqlite:///{tmp_path / 'backend.db'}"
    artifact_root = (tmp_path / "artifacts").as_uri()
    config_path = _write_trace_archival_config(
        tmp_path,
        long_retention_allowlist=["1", "2"],
        interval_seconds=42,
        max_traces_per_pass=1000,
    )

    MLFLOW_TRACE_ARCHIVAL_CONFIG.unset()

    try:
        with (
            mock.patch("mlflow.server._run_server") as run_server_mock,
            mock.patch("mlflow.server.handlers.initialize_backend_stores") as init_backend,
        ):
            result = CliRunner().invoke(
                server,
                [
                    "--backend-store-uri",
                    backend_uri,
                    "--registry-store-uri",
                    backend_uri,
                    "--default-artifact-root",
                    artifact_root,
                    "--trace-archival-config",
                    str(config_path),
                ],
                catch_exceptions=False,
                standalone_mode=False,
            )
        assert result.exit_code == 0
        run_server_mock.assert_called_once()
        init_backend.assert_called_once()
        assert MLFLOW_TRACE_ARCHIVAL_CONFIG.get() == str(config_path)
    finally:
        MLFLOW_TRACE_ARCHIVAL_CONFIG.unset()


def test_server_trace_archival_config_env_var_sets_env(tmp_path):
    handlers._tracking_store = None
    handlers._model_registry_store = None
    backend_uri = f"sqlite:///{tmp_path / 'backend.db'}"
    artifact_root = (tmp_path / "artifacts").as_uri()
    config_path = _write_trace_archival_config(tmp_path, long_retention_allowlist=["1", "2", "3"])

    MLFLOW_TRACE_ARCHIVAL_CONFIG.unset()

    try:
        MLFLOW_TRACE_ARCHIVAL_CONFIG.set(str(config_path))
        with (
            mock.patch("mlflow.server._run_server") as run_server_mock,
            mock.patch("mlflow.server.handlers.initialize_backend_stores") as init_backend,
        ):
            result = CliRunner().invoke(
                server,
                [
                    "--backend-store-uri",
                    backend_uri,
                    "--registry-store-uri",
                    backend_uri,
                    "--default-artifact-root",
                    artifact_root,
                ],
                catch_exceptions=False,
                standalone_mode=False,
            )
        assert result.exit_code == 0
        run_server_mock.assert_called_once()
        init_backend.assert_called_once()
        assert MLFLOW_TRACE_ARCHIVAL_CONFIG.get() == str(config_path)
    finally:
        MLFLOW_TRACE_ARCHIVAL_CONFIG.unset()


def test_server_trace_archival_config_cli_overrides_env_var(tmp_path):
    handlers._tracking_store = None
    handlers._model_registry_store = None
    backend_uri = f"sqlite:///{tmp_path / 'backend.db'}"
    artifact_root = (tmp_path / "artifacts").as_uri()
    env_config_path = _write_trace_archival_config(tmp_path / "env")
    cli_config_path = _write_trace_archival_config(tmp_path / "cli", interval_seconds=42)

    MLFLOW_TRACE_ARCHIVAL_CONFIG.unset()

    try:
        MLFLOW_TRACE_ARCHIVAL_CONFIG.set(str(env_config_path))
        with (
            mock.patch("mlflow.server._run_server") as run_server_mock,
            mock.patch("mlflow.server.handlers.initialize_backend_stores") as init_backend,
        ):
            result = CliRunner().invoke(
                server,
                [
                    "--backend-store-uri",
                    backend_uri,
                    "--registry-store-uri",
                    backend_uri,
                    "--default-artifact-root",
                    artifact_root,
                    "--trace-archival-config",
                    str(cli_config_path),
                ],
                catch_exceptions=False,
                standalone_mode=False,
            )
        assert result.exit_code == 0
        run_server_mock.assert_called_once()
        init_backend.assert_called_once()
        assert MLFLOW_TRACE_ARCHIVAL_CONFIG.get() == str(cli_config_path)
    finally:
        MLFLOW_TRACE_ARCHIVAL_CONFIG.unset()


def test_server_trace_archival_scheduler_defaults_disabled(tmp_path):
    handlers._tracking_store = None
    handlers._model_registry_store = None
    backend_uri = f"sqlite:///{tmp_path / 'backend.db'}"
    artifact_root = (tmp_path / "artifacts").as_uri()

    MLFLOW_TRACE_ARCHIVAL_CONFIG.unset()

    try:
        with (
            mock.patch("mlflow.server._run_server") as run_server_mock,
            mock.patch("mlflow.server.handlers.initialize_backend_stores") as init_backend,
        ):
            result = CliRunner().invoke(
                server,
                [
                    "--backend-store-uri",
                    backend_uri,
                    "--registry-store-uri",
                    backend_uri,
                    "--default-artifact-root",
                    artifact_root,
                ],
                catch_exceptions=False,
                standalone_mode=False,
            )
        assert result.exit_code == 0
        run_server_mock.assert_called_once()
        init_backend.assert_called_once()
    finally:
        MLFLOW_TRACE_ARCHIVAL_CONFIG.unset()


def test_server_trace_archival_config_missing_file(tmp_path):
    result = CliRunner().invoke(
        server,
        ["--trace-archival-config", str(tmp_path / "missing.yaml")],
    )

    assert result.exit_code != 0
    assert "does not exist" in result.output


def test_server_trace_archival_config_rejects_invalid_yaml(tmp_path):
    config_path = tmp_path / "trace-archival.yaml"
    config_path.write_text("trace_archival: [\n", encoding="utf-8")

    result = CliRunner().invoke(server, ["--trace-archival-config", str(config_path)])

    assert result.exit_code != 0
    assert "Invalid trace archival config file" in result.output
    assert "Failed to parse YAML" in result.output


def test_server_trace_archival_config_rejects_invalid_experiment_id(tmp_path):
    config_path = _write_trace_archival_config(
        tmp_path,
        long_retention_allowlist=["1", "invalid id"],
    )

    result = CliRunner().invoke(server, ["--trace-archival-config", str(config_path)])

    assert result.exit_code != 0
    assert "Invalid experiment ID: 'invalid id'" in result.output


def test_server_trace_archival_config_rejects_invalid_interval(tmp_path):
    config_path = _write_trace_archival_config(tmp_path, interval_seconds=0)

    result = CliRunner().invoke(server, ["--trace-archival-config", str(config_path)])

    assert result.exit_code != 0
    assert "'trace_archival.interval_seconds' must be a positive integer" in result.output


def test_server_trace_archival_config_rejects_invalid_max_traces_per_pass(tmp_path):
    config_path = _write_trace_archival_config(tmp_path, max_traces_per_pass=0)

    result = CliRunner().invoke(server, ["--trace-archival-config", str(config_path)])

    assert result.exit_code != 0
    assert "'trace_archival.max_traces_per_pass' must be a positive integer" in result.output


def test_server_trace_archival_config_requires_archival_capable_tracking_store(tmp_path):
    handlers._tracking_store = None
    handlers._model_registry_store = None
    backend_uri = f"sqlite:///{tmp_path / 'backend.db'}"
    artifact_root = (tmp_path / "artifacts").as_uri()
    config_path = _write_trace_archival_config(tmp_path)
    mock_tracking_store = mock.Mock(supports_trace_archival=False)

    with (
        mock.patch("mlflow.server._run_server") as run_server_mock,
        mock.patch("mlflow.server.handlers._get_tracking_store", return_value=mock_tracking_store),
        mock.patch("mlflow.server.handlers._get_model_registry_store", return_value=None),
    ):
        result = CliRunner().invoke(
            server,
            [
                "--backend-store-uri",
                backend_uri,
                "--registry-store-uri",
                backend_uri,
                "--default-artifact-root",
                artifact_root,
                "--trace-archival-config",
                str(config_path),
            ],
        )

    assert result.exit_code != 0
    assert "Error initializing backend store" in result.output
    run_server_mock.assert_not_called()


def test_server_trace_archival_config_allows_disabled_config_on_unsupported_tracking_store(
    tmp_path,
):
    handlers._tracking_store = None
    handlers._model_registry_store = None
    backend_uri = f"sqlite:///{tmp_path / 'backend.db'}"
    artifact_root = (tmp_path / "artifacts").as_uri()
    config_path = _write_trace_archival_config(tmp_path, enabled=False)
    mock_tracking_store = mock.Mock(supports_trace_archival=False)
    MLFLOW_TRACE_ARCHIVAL_CONFIG.unset()

    try:
        with (
            mock.patch("mlflow.server._run_server") as run_server_mock,
            mock.patch(
                "mlflow.server.handlers._get_tracking_store", return_value=mock_tracking_store
            ),
            mock.patch("mlflow.server.handlers._get_model_registry_store", return_value=None),
        ):
            result = CliRunner().invoke(
                server,
                [
                    "--backend-store-uri",
                    backend_uri,
                    "--registry-store-uri",
                    backend_uri,
                    "--default-artifact-root",
                    artifact_root,
                    "--trace-archival-config",
                    str(config_path),
                ],
            )

        assert result.exit_code == 0
        run_server_mock.assert_called_once()
    finally:
        MLFLOW_TRACE_ARCHIVAL_CONFIG.unset()


@pytest.mark.parametrize("command", [server])
def test_tracking_uri_validation_failure(command):
    handlers._tracking_store = None
    with mock.patch("mlflow.server._run_server") as run_server_mock:
        # SQLAlchemy expects postgresql:// not postgres://
        CliRunner().invoke(
            command,
            [
                "--backend-store-uri",
                "postgres://user:pwd@host:5432/mydb",
                "--default-artifact-root",
                "./mlruns",
            ],
        )
        run_server_mock.assert_not_called()


@pytest.mark.parametrize("command", [server])
def test_tracking_uri_validation_sql_driver_uris(command):
    handlers._tracking_store = None
    handlers._model_registry_store = None
    with (
        mock.patch("mlflow.server._run_server") as run_server_mock,
        mock.patch("mlflow.store.tracking.sqlalchemy_store.SqlAlchemyStore"),
        mock.patch("mlflow.store.model_registry.sqlalchemy_store.SqlAlchemyStore"),
    ):
        result = CliRunner().invoke(
            command,
            [
                "--backend-store-uri",
                "mysql+pymysql://user:pwd@host:5432/mydb",
                "--default-artifact-root",
                "./mlruns",
            ],
        )
        assert result.exit_code == 0
        run_server_mock.assert_called()
    # Clean up the global variables set by the server
    mlflow.set_tracking_uri(None)
    mlflow.set_registry_uri(None)


@pytest.mark.parametrize("command", [server])
def test_registry_store_uri_different_from_tracking_store(command):
    handlers._tracking_store = None
    handlers._model_registry_store = None

    from mlflow.server.handlers import (
        ModelRegistryStoreRegistryWrapper,
        TrackingStoreRegistryWrapper,
    )

    handlers._tracking_store_registry = TrackingStoreRegistryWrapper()
    handlers._model_registry_store_registry = ModelRegistryStoreRegistryWrapper()

    with (
        mock.patch("mlflow.server._run_server") as run_server_mock,
        mock.patch("mlflow.store.tracking.file_store.FileStore") as tracking_store,
        mock.patch(
            "mlflow.store.model_registry.sqlalchemy_store.SqlAlchemyStore"
        ) as registry_store,
    ):
        result = CliRunner().invoke(
            command,
            [
                "--backend-store-uri",
                "./mlruns",
                "--registry-store-uri",
                "mysql://user:pwd@host:5432/mydb",
            ],
        )
        assert result.exit_code == 0
        run_server_mock.assert_called()
        tracking_store.assert_called()
        registry_store.assert_called()
    # Clean up the global variables set by the server
    mlflow.set_tracking_uri(None)
    mlflow.set_registry_uri(None)


@pytest.fixture
def sqlite_store(db_uri: str, tmp_path: Path) -> tuple[SqlAlchemyStore, str]:
    artifact_uri = (tmp_path / "artifacts").as_uri()
    store = SqlAlchemyStore(db_uri, artifact_uri)
    return (store, db_uri)


@pytest.fixture
def file_store():
    pytest.skip("FileStore is no longer supported.")
    ROOT_LOCATION = os.path.join(tempfile.gettempdir(), "test_mlflow_gc")
    file_store_uri = f"file:///{ROOT_LOCATION}"
    yield (FileStore(ROOT_LOCATION), file_store_uri)
    shutil.rmtree(ROOT_LOCATION)


def _create_run_in_store(store, create_artifacts=True):
    config = {
        "experiment_id": "0",
        "user_id": "Anderson",
        "start_time": get_current_time_millis(),
        "tags": [],
        "run_name": "name",
    }
    run = store.create_run(**config)
    if create_artifacts:
        artifact_path = url2pathname(unquote(urlparse(run.info.artifact_uri).path))
        if not os.path.exists(artifact_path):
            os.makedirs(artifact_path)
    return run


@pytest.mark.parametrize("create_artifacts_in_run", [True, False])
def test_mlflow_gc_sqlite(sqlite_store, create_artifacts_in_run):
    store = sqlite_store[0]
    run = _create_run_in_store(store, create_artifacts=create_artifacts_in_run)
    store.delete_run(run.info.run_id)
    subprocess.check_call([
        sys.executable,
        "-m",
        "mlflow",
        "gc",
        "--backend-store-uri",
        sqlite_store[1],
    ])
    runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL)
    assert len(runs) == 0
    with pytest.raises(MlflowException, match=r"Run .+ not found"):
        store.get_run(run.info.run_id)

    artifact_path = url2pathname(unquote(urlparse(run.info.artifact_uri).path))
    assert not os.path.exists(artifact_path)


def test_mlflow_gc_sqlite_older_than(sqlite_store):
    store = sqlite_store[0]
    run = _create_run_in_store(store)
    store.delete_run(run.info.run_id)
    with pytest.raises(subprocess.CalledProcessError, match=r".+") as exp:
        subprocess.run(
            [
                sys.executable,
                "-m",
                "mlflow",
                "gc",
                "--backend-store-uri",
                sqlite_store[1],
                "--older-than",
                "10d10h10m10s",
                "--run-ids",
                run.info.run_id,
            ],
            check=True,
            capture_output=True,
            text=True,
        )
    assert "is not older than the required age" in exp.value.stderr
    runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL)
    assert len(runs) == 1

    time.sleep(1)
    subprocess.check_output([
        sys.executable,
        "-m",
        "mlflow",
        "gc",
        "--backend-store-uri",
        sqlite_store[1],
        "--older-than",
        "1s",
        "--run-ids",
        run.info.run_id,
    ])
    runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL)
    assert len(runs) == 0


@pytest.mark.parametrize("create_artifacts_in_run", [True, False])
def test_mlflow_gc_file_store(file_store, create_artifacts_in_run):
    pytest.skip("FileStore is no longer supported.")
    store = file_store[0]
    run = _create_run_in_store(store, create_artifacts=create_artifacts_in_run)
    store.delete_run(run.info.run_id)
    subprocess.check_output([
        sys.executable,
        "-m",
        "mlflow",
        "gc",
        "--backend-store-uri",
        file_store[1],
    ])
    runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL)
    assert len(runs) == 0
    with pytest.raises(MlflowException, match=r"Run .+ not found"):
        store.get_run(run.info.run_id)

    artifact_path = url2pathname(unquote(urlparse(run.info.artifact_uri).path))
    assert not os.path.exists(artifact_path)


def test_mlflow_gc_file_store_passing_explicit_run_ids(file_store):
    pytest.skip("FileStore is no longer supported.")
    store = file_store[0]
    run = _create_run_in_store(store)
    store.delete_run(run.info.run_id)
    subprocess.check_output([
        sys.executable,
        "-m",
        "mlflow",
        "gc",
        "--backend-store-uri",
        file_store[1],
        "--run-ids",
        run.info.run_id,
    ])
    runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL)
    assert len(runs) == 0
    with pytest.raises(MlflowException, match=r"Run .+ not found"):
        store.get_run(run.info.run_id)


def test_mlflow_gc_not_deleted_run(file_store):
    pytest.skip("FileStore is no longer supported.")
    store = file_store[0]
    run = _create_run_in_store(store)
    with pytest.raises(subprocess.CalledProcessError, match=r".+"):
        subprocess.check_output([
            sys.executable,
            "-m",
            "mlflow",
            "gc",
            "--backend-store-uri",
            file_store[1],
            "--run-ids",
            run.info.run_id,
        ])
    runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL)
    assert len(runs) == 1


def test_mlflow_gc_file_store_older_than(file_store):
    pytest.skip("FileStore is no longer supported.")
    store = file_store[0]
    run = _create_run_in_store(store)
    store.delete_run(run.info.run_id)
    with pytest.raises(subprocess.CalledProcessError, match=r".+") as exp:
        subprocess.run(
            [
                sys.executable,
                "-m",
                "mlflow",
                "gc",
                "--backend-store-uri",
                file_store[1],
                "--older-than",
                "10d10h10m10s",
                "--run-ids",
                run.info.run_id,
            ],
            check=True,
            capture_output=True,
            text=True,
        )
    assert "is not older than the required age" in exp.value.stderr
    runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL)
    assert len(runs) == 1

    time.sleep(1)
    subprocess.check_output([
        sys.executable,
        "-m",
        "mlflow",
        "gc",
        "--backend-store-uri",
        file_store[1],
        "--older-than",
        "1s",
        "--run-ids",
        run.info.run_id,
    ])
    runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL)
    assert len(runs) == 0


@pytest.mark.parametrize("get_store_details", ["file_store", "sqlite_store"])
def test_mlflow_gc_experiments(get_store_details, request):
    if get_store_details == "file_store":
        pytest.skip("FileStore is no longer supported.")

    def invoke_gc(*args):
        return CliRunner().invoke(gc, args, catch_exceptions=False)

    store, uri = request.getfixturevalue(get_store_details)
    exp_id_1 = store.create_experiment("1")
    run_id_1 = store.create_run(exp_id_1, user_id="user", start_time=0, tags=[], run_name="1")
    invoke_gc("--backend-store-uri", uri)
    experiments = store.search_experiments(view_type=ViewType.ALL)
    exp_ids = [e.experiment_id for e in experiments]
    runs = store.search_runs(experiment_ids=exp_ids, filter_string="", run_view_type=ViewType.ALL)
    assert sorted(exp_ids) == sorted([exp_id_1, store.DEFAULT_EXPERIMENT_ID])
    assert [r.info.run_id for r in runs] == [run_id_1.info.run_id]

    store.delete_experiment(exp_id_1)
    invoke_gc("--backend-store-uri", uri)
    experiments = store.search_experiments(view_type=ViewType.ALL)
    runs = store.search_runs(experiment_ids=exp_ids, filter_string="", run_view_type=ViewType.ALL)
    assert [e.experiment_id for e in experiments] == [store.DEFAULT_EXPERIMENT_ID]
    assert runs == []

    exp_id_2 = store.create_experiment("2")
    exp_id_3 = store.create_experiment("3")
    store.delete_experiment(exp_id_2)
    store.delete_experiment(exp_id_3)
    invoke_gc("--backend-store-uri", uri, "--experiment-ids", exp_id_2)
    experiments = store.search_experiments(view_type=ViewType.ALL)
    assert sorted([e.experiment_id for e in experiments]) == sorted([
        exp_id_3,
        store.DEFAULT_EXPERIMENT_ID,
    ])

    with mock.patch("time.time", return_value=0) as mock_time:
        exp_id_4 = store.create_experiment("4")
        store.delete_experiment(exp_id_4)
        mock_time.assert_called()

    invoke_gc("--backend-store-uri", uri, "--older-than", "1d")
    experiments = store.search_experiments(view_type=ViewType.ALL)
    assert sorted([e.experiment_id for e in experiments]) == sorted([
        exp_id_3,
        store.DEFAULT_EXPERIMENT_ID,
    ])

    invoke_gc("--backend-store-uri", uri, "--experiment-ids", exp_id_3, "--older-than", "0s")
    experiments = store.search_experiments(view_type=ViewType.ALL)
    assert [e.experiment_id for e in experiments] == [store.DEFAULT_EXPERIMENT_ID]

    exp_id_5 = store.create_experiment("5")
    store.delete_experiment(exp_id_5)
    with pytest.raises(MlflowException, match=r"Experiments .+ can be deleted."):
        invoke_gc(
            "--backend-store-uri",
            uri,
            "--experiment-ids",
            exp_id_5,
            "--older-than",
            "10d10h10m10s",
        )
    experiments = store.search_experiments(view_type=ViewType.ALL)
    assert sorted([e.experiment_id for e in experiments]) == sorted([
        exp_id_5,
        store.DEFAULT_EXPERIMENT_ID,
    ])


def test_mlflow_gc_experiment_with_logged_model_params_tags_and_metrics(sqlite_store):
    store, uri = sqlite_store
    exp_id = store.create_experiment("exp_with_logged_model")
    run = store.create_run(exp_id, user_id="user", start_time=0, tags=[], run_name="run")
    run_id = run.info.run_id
    model = store.create_logged_model(experiment_id=exp_id)

    store.log_logged_model_params(
        model_id=model.model_id,
        params=[
            LoggedModelParameter(key="param1", value="value1"),
            LoggedModelParameter(key="param2", value="value2"),
        ],
    )
    store.set_logged_model_tags(
        model_id=model.model_id,
        tags=[
            LoggedModelTag(key="tag1", value="value1"),
            LoggedModelTag(key="tag2", value="value2"),
        ],
    )
    store._log_model_metrics(
        run_id=run_id,
        metrics=[Metric(key="m1", value=1.0, timestamp=0, step=0, model_id=model.model_id)],
        experiment_id=exp_id,
    )

    store.delete_experiment(exp_id)

    result = CliRunner().invoke(gc, ["--backend-store-uri", uri], catch_exceptions=False)
    assert result.exit_code == 0

    experiments = store.search_experiments(view_type=ViewType.ALL)
    assert [e.experiment_id for e in experiments] == [store.DEFAULT_EXPERIMENT_ID]

    with pytest.raises(MlflowException, match=r".+ not found"):
        store.get_logged_model(model.model_id)

    with store.ManagedSessionMaker() as session:
        for table in [SqlLoggedModelParam, SqlLoggedModelTag, SqlLoggedModelMetric]:
            assert session.query(table).count() == 0


@pytest.fixture
def sqlite_store_with_s3_artifact_repository(
    tmp_path: Path,
) -> tuple[SqlAlchemyStore, str, str]:
    db_path = tmp_path / "mlflow.db"
    db_uri = f"sqlite:///{db_path}"
    s3_uri = "s3://mlflow"
    store = SqlAlchemyStore(db_uri, s3_uri)
    return (store, db_uri, s3_uri)


def test_mlflow_gc_sqlite_with_s3_artifact_repository(
    sqlite_store_with_s3_artifact_repository,
):
    store = sqlite_store_with_s3_artifact_repository[0]
    run = _create_run_in_store(store, create_artifacts=False)
    store.delete_run(run.info.run_id)

    artifact_repo = get_artifact_repository(run.info.artifact_uri)
    bucket, dest_path = artifact_repo.parse_s3_compliant_uri(run.info.artifact_uri)
    fake_artifact_path = os.path.join(dest_path, "fake_artifact.txt")
    with Stubber(artifact_repo._get_s3_client()) as s3_stubber:
        s3_stubber.add_response(
            "list_objects_v2",
            {"Contents": [{"Key": fake_artifact_path}]},
            {"Bucket": bucket, "Prefix": dest_path},
        )
        s3_stubber.add_response(
            "delete_objects",
            {"Deleted": [{"Key": fake_artifact_path}]},
            {"Bucket": bucket, "Delete": {"Objects": [{"Key": fake_artifact_path}]}},
        )

        CliRunner().invoke(
            gc,
            [
                "--backend-store-uri",
                sqlite_store_with_s3_artifact_repository[1],
                "--artifacts-destination",
                sqlite_store_with_s3_artifact_repository[2],
            ],
            catch_exceptions=False,
        )

        runs = store.search_runs(experiment_ids=["0"], filter_string="", run_view_type=ViewType.ALL)
        assert len(runs) == 0
        with pytest.raises(MlflowException, match=r"Run .+ not found"):
            store.get_run(run.info.run_id)


def test_mlflow_tracking_disabled_in_artifacts_only_mode(tmp_path: Path):
    port = get_safe_port()
    env = {**os.environ}
    env.pop(MLFLOW_TRACE_ARCHIVAL_CONFIG.name, None)
    with subprocess.Popen(
        [sys.executable, "-m", "mlflow", "server", "--port", str(port), "--artifacts-only"],
        cwd=tmp_path,
        env=env,
    ) as process:
        try:
            _await_server_up_or_die(port)
            resp = requests.get(f"http://localhost:{port}/api/2.0/mlflow/experiments/search")
            assert (
                "Endpoint: /api/2.0/mlflow/experiments/search disabled due to the mlflow "
                "server running in `--artifacts-only` mode." in resp.text
            )
        finally:
            kill_process_tree(process.pid)


def test_mlflow_artifact_list_in_artifacts_only_mode(tmp_path: Path):
    port = get_safe_port()
    env = {**os.environ}
    env.pop(MLFLOW_TRACE_ARCHIVAL_CONFIG.name, None)
    with subprocess.Popen(
        [sys.executable, "-m", "mlflow", "server", "--port", str(port), "--artifacts-only"],
        cwd=tmp_path,
        env=env,
    ) as process:
        try:
            _await_server_up_or_die(port)
            resp = requests.get(f"http://localhost:{port}/api/2.0/mlflow-artifacts/artifacts")
            augmented_raise_for_status(resp)
            assert resp.status_code == 200
            assert resp.text == "{}"
        finally:
            kill_process_tree(process.pid)


def test_mlflow_artifact_service_unavailable_when_no_server_artifacts_is_specified():
    port = get_safe_port()
    with subprocess.Popen([
        sys.executable,
        "-m",
        "mlflow",
        "server",
        "--port",
        str(port),
        "--no-serve-artifacts",
    ]) as process:
        try:
            _await_server_up_or_die(port)
            endpoint = "/api/2.0/mlflow-artifacts/artifacts"
            resp = requests.get(f"http://localhost:{port}{endpoint}")
            assert (
                f"Endpoint: {endpoint} disabled due to the mlflow server running with "
                "`--no-serve-artifacts`" in resp.text
            )
        finally:
            kill_process_tree(process.pid)


def test_mlflow_artifact_only_prints_warning_for_configs():
    with (
        mock.patch("mlflow.server._run_server") as run_server_mock,
        mock.patch("mlflow.store.tracking.sqlalchemy_store.SqlAlchemyStore"),
        mock.patch("mlflow.store.model_registry.sqlalchemy_store.SqlAlchemyStore"),
    ):
        result = CliRunner().invoke(
            server,
            ["--artifacts-only", "--backend-store-uri", "sqlite:///my.db"],
            catch_exceptions=False,
        )
        msg = (
            "Usage: server [OPTIONS]\nTry 'server --help' for help.\n\nError: You are starting a "
            "tracking server in `--artifacts-only` mode and have provided a value for "
            "`--backend_store_uri`"
        )
        assert msg in result.output
        assert result.exit_code != 0
        run_server_mock.assert_not_called()


def test_mlflow_ui_is_alias_for_mlflow_server():
    mlflow_ui_stdout = subprocess.check_output(
        [sys.executable, "-m", "mlflow", "ui", "--help"], text=True
    )
    mlflow_server_stdout = subprocess.check_output(
        [sys.executable, "-m", "mlflow", "server", "--help"], text=True
    )
    assert (
        mlflow_ui_stdout.replace("Usage: python -m mlflow ui", "Usage: python -m mlflow server")
        == mlflow_server_stdout
    )


def test_server_help_has_quickstart():
    out = subprocess.check_output([sys.executable, "-m", "mlflow", "server", "--help"], text=True)
    assert "Pick your setup" in out
    assert "Options by topic:" in out
    assert "sqlite:///mlflow.db" in out
    # Curated guidance is rendered before the exhaustive flat option list.
    assert out.index("Pick your setup") < out.index("\nOptions:")


def test_cli_with_python_mod():
    stdout = subprocess.check_output(
        [
            sys.executable,
            "-m",
            "mlflow",
            "--version",
        ],
        text=True,
    )
    assert stdout.rstrip().endswith(mlflow.__version__)
    stdout = subprocess.check_output(
        [
            sys.executable,
            "-m",
            "mlflow",
            "server",
            "--help",
        ],
        text=True,
    )
    assert "mlflow server" in stdout


def test_doctor():
    res = CliRunner().invoke(doctor, catch_exceptions=False)
    assert res.exit_code == 0


def test_env_file_loading(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
    # Setup: Create an experiment using the Python SDK
    test_tracking_uri = f"sqlite:///{tmp_path / 'mlflow.db'}"
    test_experiment_name = "test_experiment_from_env"

    # Create experiment using SDK
    mlflow.set_tracking_uri(test_tracking_uri)
    mlflow.create_experiment(test_experiment_name)

    # Create a test .env file pointing to this tracking URI
    env_file_path = tmp_path / "test.env"
    env_file_path.write_text(f"MLFLOW_TRACKING_URI={test_tracking_uri}\n")

    runner = CliRunner()

    # Clear the MLflow environment variables to ensure we're testing the loading
    monkeypatch.delenv("MLFLOW_TRACKING_URI", raising=False)
    monkeypatch.delenv("MLFLOW_EXPERIMENT_NAME", raising=False)

    # Ensure variables are not set before running command
    assert "MLFLOW_TRACKING_URI" not in os.environ

    # Use the existing experiments search CLI command with --env-file
    result = runner.invoke(
        cli,
        ["--env-file", str(env_file_path), "experiments", "search"],
        catch_exceptions=False,
    )

    # Check that the command executed successfully
    assert result.exit_code == 0

    # Verify the experiment we created is found (proves env vars were loaded)
    assert test_experiment_name in result.output

    # Verify the loading message
    assert "Loaded environment variables from:" in result.output
    assert str(env_file_path) in result.output


def test_env_file_loading_invalid_path() -> None:
    runner = CliRunner()

    # Test error handling for non-existent file
    result = runner.invoke(
        cli,
        ["--env-file", "nonexistent.env", "experiments", "search"],
        catch_exceptions=False,
    )
    assert result.exit_code != 0
    assert "Environment file 'nonexistent.env' does not exist" in result.output


@pytest.mark.parametrize(
    ("args", "expected_max_results", "expected_view"),
    [
        (["experiments", "search"], None, ViewType.ACTIVE_ONLY),
        (["experiments", "search", "--max-results", "50"], 50, ViewType.ACTIVE_ONLY),
        (
            ["experiments", "search", "--view", "all", "--max-results", "100"],
            100,
            ViewType.ALL,
        ),
    ],
)
def test_experiments_search_max_results(args, expected_max_results, expected_view):
    with mock.patch("mlflow.search_experiments", return_value=[]) as mock_search:
        result = CliRunner().invoke(cli, args, catch_exceptions=False)
        assert result.exit_code == 0
        mock_search.assert_called_once()
        call_kwargs = mock_search.call_args[1]
        assert call_kwargs.get("max_results") == expected_max_results
        assert call_kwargs.get("view_type") == expected_view


def test_experiments_search_max_results_zero():
    with mock.patch("mlflow.search_experiments", return_value=[]) as mock_search:
        result = CliRunner().invoke(
            cli, ["experiments", "search", "--max-results", "0"], catch_exceptions=False
        )
        assert result.exit_code == 0
        mock_search.assert_called_once()
        call_kwargs = mock_search.call_args[1]
        assert call_kwargs.get("max_results") == 0


def test_experiments_search_max_results_negative():
    result = CliRunner().invoke(
        cli, ["experiments", "search", "--max-results", "-1"], catch_exceptions=False
    )
    assert result.exit_code != 0
    assert "max-results must be a non-negative integer" in result.output


def test_server_with_env_file(tmp_path):
    env_file = tmp_path / ".env"
    env_file.write_text("TEST_VAR=test_value\n")

    with mock.patch("mlflow.server._run_server") as run_server_mock:
        result = CliRunner().invoke(cli, ["--env-file", str(env_file), "server"])
        assert result.exit_code == 0
        run_server_mock.assert_called_once()
        # Verify env_file parameter is passed
        assert run_server_mock.call_args.kwargs["env_file"] == str(env_file)


def test_mlflow_gc_with_datasets(sqlite_store):
    store = sqlite_store[0]

    mlflow.set_tracking_uri(sqlite_store[1])
    mlflow.set_experiment("dataset")

    dataset = numpy_dataset.from_numpy(np.array([1, 2, 3]))

    with mlflow.start_run() as run:
        experiment_id = run.info.experiment_id
        mlflow.log_input(dataset)

    experiments = store.search_experiments(view_type=ViewType.ALL)

    # default and datasets
    assert len(experiments) == 2

    store.delete_experiment(experiment_id)

    # the new experiment is only marked as deleted, not removed
    experiments = store.search_experiments(view_type=ViewType.ALL)
    assert len(experiments) == 2

    subprocess.check_call([
        sys.executable,
        "-m",
        "mlflow",
        "gc",
        "--backend-store-uri",
        sqlite_store[1],
    ])
    experiments = store.search_experiments(view_type=ViewType.ALL)

    # only default is left after GC
    assert len(experiments) == 1
    assert experiments[0].experiment_id == "0"
    with pytest.raises(MlflowException, match=f"No Experiment with id={experiment_id} exists"):
        store.get_experiment(experiment_id)


@pytest.mark.parametrize("get_store_details", ["file_store", "sqlite_store"])
def test_mlflow_gc_logged_model(get_store_details, request):
    if get_store_details == "file_store":
        pytest.skip("FileStore is no longer supported.")
    store, uri = request.getfixturevalue(get_store_details)
    exp_id = store.create_experiment("exp")
    model = store.create_logged_model(experiment_id=exp_id)

    store.delete_logged_model(model.model_id)
    subprocess.check_output([sys.executable, "-m", "mlflow", "gc", "--backend-store-uri", uri])

    with pytest.raises(MlflowException, match=r".+ not found"):
        store.get_logged_model(model.model_id)


@pytest.mark.parametrize("get_store_details", ["file_store", "sqlite_store"])
def test_mlflow_gc_logged_models_older_than(get_store_details, request):
    if get_store_details == "file_store":
        pytest.skip("FileStore is no longer supported.")
    store, uri = request.getfixturevalue(get_store_details)
    exp_id = store.create_experiment("exp")
    old_time = time.time() - (2 * 24 * 60 * 60)
    with mock.patch("time.time", return_value=old_time):
        model = store.create_logged_model(experiment_id=exp_id)

    store.delete_logged_model(model.model_id)

    subprocess.check_call([
        sys.executable,
        "-m",
        "mlflow",
        "gc",
        "--backend-store-uri",
        uri,
        "--older-than",
        "1d",
    ])

    retrieved_model = store.get_logged_model(model.model_id, allow_deleted=True)
    assert retrieved_model.model_id == model.model_id


@pytest.mark.parametrize("get_store_details", ["file_store", "sqlite_store"])
def test_mlflow_gc_logged_models_deletes_when_older_than(get_store_details, request):
    if get_store_details == "file_store":
        pytest.skip("FileStore is no longer supported.")
    store, uri = request.getfixturevalue(get_store_details)
    exp_id = store.create_experiment("exp")
    model = store.create_logged_model(experiment_id=exp_id)

    old_deletion_ms = int((time.time() - (2 * 24 * 60 * 60)) * 1000)
    with mock.patch("mlflow.utils.time.get_current_time_millis", return_value=old_deletion_ms):
        store.delete_logged_model(model.model_id)

    subprocess.check_call([
        sys.executable,
        "-m",
        "mlflow",
        "gc",
        "--backend-store-uri",
        uri,
        "--older-than",
        "1d",
    ])

    with pytest.raises(MlflowException, match=r".+ not found"):
        store.get_logged_model(model.model_id)


@pytest.mark.parametrize("get_store_details", ["file_store", "sqlite_store"])
def test_mlflow_gc_logged_models_mixed_time(get_store_details, request):
    if get_store_details == "file_store":
        pytest.skip("FileStore is no longer supported.")
    store, uri = request.getfixturevalue(get_store_details)
    exp_id = store.create_experiment("exp")
    old_model = store.create_logged_model(experiment_id=exp_id)
    recent_model = store.create_logged_model(experiment_id=exp_id)

    old_deletion_ms = int((time.time() - (3 * 24 * 60 * 60)) * 1000)
    with mock.patch("mlflow.utils.time.get_current_time_millis", return_value=old_deletion_ms):
        store.delete_logged_model(old_model.model_id)

    with mock.patch("mlflow.utils.time.get_current_time_millis") as current_time_mock:
        current_time_mock.return_value = int(time.time() * 1000)
        store.delete_logged_model(recent_model.model_id)

    subprocess.check_call([
        sys.executable,
        "-m",
        "mlflow",
        "gc",
        "--backend-store-uri",
        uri,
        "--older-than",
        "1d",
    ])

    with pytest.raises(MlflowException, match=r".+ not found"):
        store.get_logged_model(old_model.model_id)

    retrieved_model = store.get_logged_model(recent_model.model_id, allow_deleted=True)
    assert retrieved_model.model_id == recent_model.model_id


@pytest.fixture
def sqlite_store_with_jobs(
    db_uri: str, tmp_path: Path
) -> tuple[SqlAlchemyStore, SqlAlchemyJobStore, str]:
    artifact_uri = (tmp_path / "artifacts").as_uri()
    tracking_store = SqlAlchemyStore(db_uri, artifact_uri)
    job_store = SqlAlchemyJobStore(db_uri)
    return (tracking_store, job_store, db_uri)


def _create_test_job(job_store, job_name="test_job", finalize=True):
    job = job_store.create_job(job_name=job_name, params="{}")
    if finalize:
        job_store.start_job(job.job_id)
        job_store.finish_job(job.job_id, result="done")
        return job_store.get_job(job.job_id)
    return job


def _set_job_creation_time(job_store, job_id, creation_time_ms):
    from mlflow.store.tracking.dbmodels.models import SqlJob

    with job_store.ManagedSessionMaker(read_only=False) as session:
        job = session.query(SqlJob).filter(SqlJob.id == job_id).first()
        job.creation_time = creation_time_ms


def test_mlflow_gc_no_jobs_deleted_without_jobs_flag(sqlite_store_with_jobs):
    _, job_store, db_uri = sqlite_store_with_jobs

    job1 = _create_test_job(job_store, "job1")
    job2 = _create_test_job(job_store, "job2")

    subprocess.check_call([
        sys.executable,
        "-m",
        "mlflow",
        "gc",
        "--backend-store-uri",
        db_uri,
    ])

    retrieved_job1 = job_store.get_job(job1.job_id)
    retrieved_job2 = job_store.get_job(job2.job_id)
    assert retrieved_job1.job_id == job1.job_id
    assert retrieved_job2.job_id == job2.job_id


def test_mlflow_gc_all_jobs_deleted_with_jobs_flag(sqlite_store_with_jobs):
    _, job_store, db_uri = sqlite_store_with_jobs

    job1 = _create_test_job(job_store, "job1")
    job2 = _create_test_job(job_store, "job2")
    job3 = _create_test_job(job_store, "job3")

    subprocess.check_call([
        sys.executable,
        "-m",
        "mlflow",
        "gc",
        "--backend-store-uri",
        db_uri,
        "--jobs",
    ])

    with pytest.raises(MlflowException, match=r"Job .+ not found"):
        job_store.get_job(job1.job_id)
    with pytest.raises(MlflowException, match=r"Job .+ not found"):
        job_store.get_job(job2.job_id)
    with pytest.raises(MlflowException, match=r"Job .+ not found"):
        job_store.get_job(job3.job_id)


def test_mlflow_gc_specific_jobs_deleted_with_job_ids(sqlite_store_with_jobs):
    _, job_store, db_uri = sqlite_store_with_jobs

    job1 = _create_test_job(job_store, "job1")
    job2 = _create_test_job(job_store, "job2")
    job3 = _create_test_job(job_store, "job3")

    subprocess.check_call([
        sys.executable,
        "-m",
        "mlflow",
        "gc",
        "--backend-store-uri",
        db_uri,
        "--job-ids",
        f"{job1.job_id},{job2.job_id}",
    ])

    with pytest.raises(MlflowException, match=r"Job .+ not found"):
        job_store.get_job(job1.job_id)
    with pytest.raises(MlflowException, match=r"Job .+ not found"):
        job_store.get_job(job2.job_id)

    retrieved_job3 = job_store.get_job(job3.job_id)
    assert retrieved_job3.job_id == job3.job_id


def test_mlflow_gc_jobs_deleted_with_older_than_flag(sqlite_store_with_jobs):
    _, job_store, db_uri = sqlite_store_with_jobs

    old_job = _create_test_job(job_store, "old_job")
    new_job = _create_test_job(job_store, "new_job")

    one_hour_ago = get_current_time_millis() - (60 * 60 * 1000)
    _set_job_creation_time(job_store, old_job.job_id, one_hour_ago)

    subprocess.check_call([
        sys.executable,
        "-m",
        "mlflow",
        "gc",
        "--backend-store-uri",
        db_uri,
        "--jobs",
        "--older-than",
        "30m",
    ])

    with pytest.raises(MlflowException, match=r"Job .+ not found"):
        job_store.get_job(old_job.job_id)

    retrieved_new_job = job_store.get_job(new_job.job_id)
    assert retrieved_new_job.job_id == new_job.job_id


def test_mlflow_gc_job_ids_with_older_than_filter(sqlite_store_with_jobs):
    _, job_store, db_uri = sqlite_store_with_jobs

    old_job = _create_test_job(job_store, "old_job")
    new_job = _create_test_job(job_store, "new_job")

    one_hour_ago = get_current_time_millis() - (60 * 60 * 1000)
    _set_job_creation_time(job_store, old_job.job_id, one_hour_ago)

    subprocess.check_call([
        sys.executable,
        "-m",
        "mlflow",
        "gc",
        "--backend-store-uri",
        db_uri,
        "--job-ids",
        new_job.job_id,
        "--older-than",
        "30m",
    ])

    retrieved_new_job = job_store.get_job(new_job.job_id)
    assert retrieved_new_job.job_id == new_job.job_id

    retrieved_old_job = job_store.get_job(old_job.job_id)
    assert retrieved_old_job.job_id == old_job.job_id


def test_mlflow_gc_only_deletes_finalized_jobs(sqlite_store_with_jobs):
    from mlflow.entities._job_status import JobStatus

    _, job_store, db_uri = sqlite_store_with_jobs

    # Create jobs with different statuses
    pending_job = _create_test_job(job_store, "pending_job", finalize=False)
    assert pending_job.status == JobStatus.PENDING

    running_job = _create_test_job(job_store, "running_job", finalize=False)
    job_store.start_job(running_job.job_id)
    running_job = job_store.get_job(running_job.job_id)
    assert running_job.status == JobStatus.RUNNING

    succeeded_job = _create_test_job(job_store, "succeeded_job", finalize=True)
    assert succeeded_job.status == JobStatus.SUCCEEDED

    failed_job = _create_test_job(job_store, "failed_job", finalize=False)
    job_store.start_job(failed_job.job_id)
    job_store.fail_job(failed_job.job_id, "error")
    failed_job = job_store.get_job(failed_job.job_id)
    assert failed_job.status == JobStatus.FAILED

    timeout_job = _create_test_job(job_store, "timeout_job", finalize=False)
    job_store.start_job(timeout_job.job_id)
    job_store.mark_job_timed_out(timeout_job.job_id)
    timeout_job = job_store.get_job(timeout_job.job_id)
    assert timeout_job.status == JobStatus.TIMEOUT

    canceled_job = _create_test_job(job_store, "canceled_job", finalize=False)
    job_store.cancel_job(canceled_job.job_id)
    canceled_job = job_store.get_job(canceled_job.job_id)
    assert canceled_job.status == JobStatus.CANCELED

    subprocess.check_call([
        sys.executable,
        "-m",
        "mlflow",
        "gc",
        "--backend-store-uri",
        db_uri,
        "--jobs",
    ])

    retrieved_pending = job_store.get_job(pending_job.job_id)
    assert retrieved_pending.job_id == pending_job.job_id
    assert retrieved_pending.status == JobStatus.PENDING

    retrieved_running = job_store.get_job(running_job.job_id)
    assert retrieved_running.job_id == running_job.job_id
    assert retrieved_running.status == JobStatus.RUNNING

    with pytest.raises(MlflowException, match=r"Job .+ not found"):
        job_store.get_job(succeeded_job.job_id)
    with pytest.raises(MlflowException, match=r"Job .+ not found"):
        job_store.get_job(failed_job.job_id)
    with pytest.raises(MlflowException, match=r"Job .+ not found"):
        job_store.get_job(timeout_job.job_id)
    with pytest.raises(MlflowException, match=r"Job .+ not found"):
        job_store.get_job(canceled_job.job_id)