wehub-resource-sync
415 lines
15 KiB
Python
415 lines
15 KiB
Python
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
|