mlflow--mlflow/tests/server/auth/test_read_write_routing.py

import shutil
from pathlib import Path
from unittest import mock

import pytest

from mlflow.server.auth.config import AuthConfig, read_auth_config
from mlflow.server.auth.sqlalchemy_store import SqlAlchemyStore


@pytest.fixture
def write_db_uri(tmp_path: Path) -> str:
    return f"sqlite:///{tmp_path / 'auth_write.db'}"


@pytest.fixture
def read_db_uri(tmp_path: Path) -> str:
    return f"sqlite:///{tmp_path / 'auth_read.db'}"


@pytest.fixture
def store_with_replica(tmp_path: Path, write_db_uri, read_db_uri):
    store = SqlAlchemyStore()
    store.init_db(write_db_uri)
    store.engine.dispose()
    shutil.copy(tmp_path / "auth_write.db", tmp_path / "auth_read.db")
    store2 = SqlAlchemyStore()
    store2.init_db(write_db_uri, read_db_uri=read_db_uri)
    return store2


@pytest.fixture
def store_no_replica(tmp_path: Path, write_db_uri):
    store = SqlAlchemyStore()
    store.init_db(write_db_uri)
    return store


# --- TestDualEngineSetup ---


def test_no_replica_uses_single_engine(store_no_replica):
    assert store_no_replica.read_engine is None


def test_replica_creates_separate_engine(store_with_replica):
    assert store_with_replica.read_engine is not None
    assert store_with_replica.engine is not store_with_replica.read_engine


def test_same_uri_skips_replica(tmp_path: Path, write_db_uri):
    store = SqlAlchemyStore()
    store.init_db(write_db_uri, read_db_uri=write_db_uri)
    assert store.read_engine is None


def test_none_read_uri_skips_replica(store_no_replica):
    assert store_no_replica.read_engine is None


# --- TestReadWriteRouting ---


def test_write_goes_to_primary(store_with_replica):
    user = store_with_replica.create_user("test_user", "password12345678")
    assert user.username == "test_user"


def test_read_from_replica(store_with_replica, tmp_path, write_db_uri, read_db_uri):
    read_store = SqlAlchemyStore()
    read_store.init_db(read_db_uri)
    read_store.create_user("replica_user", "password12345678")

    user = store_with_replica.get_user("replica_user")
    assert user.username == "replica_user"


def test_has_user_uses_replica(store_with_replica, read_db_uri):
    read_store = SqlAlchemyStore()
    read_store.init_db(read_db_uri)
    read_store.create_user("check_user", "password12345678")

    assert store_with_replica.has_user("check_user") is True


def test_authenticate_uses_replica(store_with_replica, read_db_uri):
    read_store = SqlAlchemyStore()
    read_store.init_db(read_db_uri)
    read_store.create_user("auth_user", "secretpass12345678")

    assert store_with_replica.authenticate_user("auth_user", "secretpass12345678") is True
    assert store_with_replica.authenticate_user("auth_user", "wrongpass12345678") is False


# --- TestBackwardCompatibility ---


def test_backward_compat_create_and_read_user(store_no_replica):
    store_no_replica.create_user("compat_user", "password12345678")
    user = store_no_replica.get_user("compat_user")
    assert user.username == "compat_user"


def test_backward_compat_list_users(store_no_replica):
    store_no_replica.create_user("list_user", "password12345678")
    users = store_no_replica.list_users()
    assert any(u.username == "list_user" for u in users)


def test_backward_compat_authenticate(store_no_replica):
    store_no_replica.create_user("auth_compat", "mypassword12345678")
    assert store_no_replica.authenticate_user("auth_compat", "mypassword12345678") is True


# --- TestAuthConfig ---


def test_config_with_read_uri(tmp_path):
    config = AuthConfig(
        default_permission="READ",
        database_uri="sqlite:///auth.db",
        read_database_uri="sqlite:///auth_read.db",
        admin_username="admin",
        admin_password="password",
        authorization_function="mlflow.server.auth:authenticate_request_basic_auth",
        grant_default_workspace_access=False,
        workspace_cache_max_size=10000,
        workspace_cache_ttl_seconds=3600,
        auth_cache_max_size=10000,
        auth_cache_ttl_seconds=3600,
    )
    assert config.read_database_uri == "sqlite:///auth_read.db"


def test_config_without_read_uri(tmp_path):
    config = AuthConfig(
        default_permission="READ",
        database_uri="sqlite:///auth.db",
        read_database_uri=None,
        admin_username="admin",
        admin_password="password",
        authorization_function="mlflow.server.auth:authenticate_request_basic_auth",
        grant_default_workspace_access=False,
        workspace_cache_max_size=10000,
        workspace_cache_ttl_seconds=3600,
        auth_cache_max_size=10000,
        auth_cache_ttl_seconds=3600,
    )
    assert config.read_database_uri is None


def test_read_auth_config_without_read_uri(tmp_path):
    ini_path = tmp_path / "test_auth.ini"
    ini_path.write_text(
        "[mlflow]\n"
        "default_permission = READ\n"
        "database_uri = sqlite:///test.db\n"
        "admin_username = admin\n"
        "admin_password = password\n"
    )

    with mock.patch("mlflow.server.auth.config._get_auth_config_path", return_value=str(ini_path)):
        config = read_auth_config()
        assert config.read_database_uri is None
        assert config.database_uri == "sqlite:///test.db"


def test_read_auth_config_with_read_uri(tmp_path):
    ini_path = tmp_path / "test_auth.ini"
    ini_path.write_text(
        "[mlflow]\n"
        "default_permission = READ\n"
        "database_uri = sqlite:///test.db\n"
        "read_database_uri = sqlite:///test_read.db\n"
        "admin_username = admin\n"
        "admin_password = password\n"
    )

    with mock.patch("mlflow.server.auth.config._get_auth_config_path", return_value=str(ini_path)):
        config = read_auth_config()
        assert config.read_database_uri == "sqlite:///test_read.db"