wehub-resource-sync
223 lines
8.3 KiB
Python
223 lines
8.3 KiB
Python
import pytest
|
|
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.server.auth.permissions import (
|
|
ALL_PERMISSIONS,
|
|
EDIT,
|
|
MANAGE,
|
|
NO_PERMISSIONS,
|
|
PERMISSION_PRIORITY,
|
|
READ,
|
|
RESOURCE_GRANTABLE_PERMISSIONS,
|
|
USE,
|
|
VALID_RESOURCE_TYPES,
|
|
WORKSPACE_GRANTABLE_PERMISSIONS,
|
|
_validate_permission,
|
|
_validate_permission_for_resource_type,
|
|
_validate_resource_type,
|
|
get_permission,
|
|
max_permission,
|
|
)
|
|
|
|
# ---- Permission hierarchy ---------------------------------------------------
|
|
|
|
# The canonical ordering the rest of the auth layer relies on.
|
|
_EXPECTED_ORDER = [NO_PERMISSIONS, READ, USE, EDIT, MANAGE]
|
|
|
|
|
|
def test_permission_priority_is_total_order():
|
|
priorities = [PERMISSION_PRIORITY[p.name] for p in _EXPECTED_ORDER]
|
|
assert priorities == sorted(priorities)
|
|
assert len(set(priorities)) == len(priorities)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("permission", "can_read", "can_use", "can_update", "can_delete", "can_manage"),
|
|
[
|
|
(NO_PERMISSIONS, False, False, False, False, False),
|
|
(READ, True, False, False, False, False),
|
|
(USE, True, True, False, False, False),
|
|
(EDIT, True, True, True, False, False),
|
|
(MANAGE, True, True, True, True, True),
|
|
],
|
|
)
|
|
def test_permission_capability_matrix(
|
|
permission, can_read, can_use, can_update, can_delete, can_manage
|
|
):
|
|
"""Each permission level exposes exactly the capabilities it should.
|
|
|
|
This pins the capability semantics: upgrading READ → USE adds can_use,
|
|
USE → EDIT adds can_update, EDIT → MANAGE adds can_delete AND can_manage.
|
|
If a capability bit shifts without the corresponding test update, this
|
|
catches it.
|
|
"""
|
|
assert permission.can_read is can_read
|
|
assert permission.can_use is can_use
|
|
assert permission.can_update is can_update
|
|
assert permission.can_delete is can_delete
|
|
assert permission.can_manage is can_manage
|
|
|
|
|
|
@pytest.mark.parametrize("permission", _EXPECTED_ORDER)
|
|
def test_get_permission_roundtrip(permission):
|
|
assert get_permission(permission.name) is permission
|
|
|
|
|
|
def test_all_permissions_dict_is_complete():
|
|
assert set(ALL_PERMISSIONS) == {p.name for p in _EXPECTED_ORDER}
|
|
|
|
|
|
# ---- max_permission ---------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("a", "b", "expected"),
|
|
[
|
|
# Identical inputs return the same value.
|
|
("READ", "READ", "READ"),
|
|
("MANAGE", "MANAGE", "MANAGE"),
|
|
# Higher on the right wins.
|
|
("READ", "USE", "USE"),
|
|
("USE", "EDIT", "EDIT"),
|
|
("EDIT", "MANAGE", "MANAGE"),
|
|
# Higher on the left wins.
|
|
("USE", "READ", "USE"),
|
|
("EDIT", "USE", "EDIT"),
|
|
("MANAGE", "EDIT", "MANAGE"),
|
|
# NO_PERMISSIONS loses to everything — a user with any active grant
|
|
# beats an explicit deny.
|
|
("NO_PERMISSIONS", "READ", "READ"),
|
|
("READ", "NO_PERMISSIONS", "READ"),
|
|
("NO_PERMISSIONS", "MANAGE", "MANAGE"),
|
|
("NO_PERMISSIONS", "NO_PERMISSIONS", "NO_PERMISSIONS"),
|
|
],
|
|
)
|
|
def test_max_permission(a, b, expected):
|
|
assert max_permission(a, b) == expected
|
|
|
|
|
|
def test_max_permission_is_symmetric():
|
|
for a_perm in _EXPECTED_ORDER:
|
|
for b_perm in _EXPECTED_ORDER:
|
|
forward = max_permission(a_perm.name, b_perm.name)
|
|
reverse = max_permission(b_perm.name, a_perm.name)
|
|
assert forward == reverse
|
|
|
|
|
|
def test_max_permission_treats_unknown_as_lowest():
|
|
# Defensive behavior: an unknown permission name should not silently win.
|
|
assert max_permission("UNKNOWN", "READ") == "READ"
|
|
assert max_permission("READ", "UNKNOWN") == "READ"
|
|
# Two unknowns → the first wins (fallback to priority 0 for both, ``>=`` breaks the tie).
|
|
assert max_permission("UNKNOWN_A", "UNKNOWN_B") == "UNKNOWN_A"
|
|
|
|
|
|
# ---- Validators -------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.parametrize("permission", [p.name for p in _EXPECTED_ORDER])
|
|
def test_validate_permission_accepts_known(permission):
|
|
_validate_permission(permission)
|
|
|
|
|
|
@pytest.mark.parametrize("bogus", ["", "read", "ADMIN", "Manage", "NONE"])
|
|
def test_validate_permission_rejects_unknown(bogus):
|
|
with pytest.raises(MlflowException, match="Invalid permission"):
|
|
_validate_permission(bogus)
|
|
|
|
|
|
@pytest.mark.parametrize("resource_type", sorted(VALID_RESOURCE_TYPES))
|
|
def test_validate_resource_type_accepts_known(resource_type):
|
|
_validate_resource_type(resource_type)
|
|
|
|
|
|
@pytest.mark.parametrize("bogus", ["", "Experiment", "workspaces", "run", "trace"])
|
|
def test_validate_resource_type_rejects_unknown(bogus):
|
|
with pytest.raises(MlflowException, match="Invalid resource type"):
|
|
_validate_resource_type(bogus)
|
|
|
|
|
|
def test_workspace_is_a_valid_resource_type():
|
|
"""Regression guard: the role-based permission model depends on
|
|
``workspace`` being a first-class resource type for the workspace-wide
|
|
grant shape ``(workspace, *, PERMISSION)``. If it's ever removed from
|
|
VALID_RESOURCE_TYPES, the UI and the role-resolution logic both break.
|
|
"""
|
|
assert "workspace" in VALID_RESOURCE_TYPES
|
|
|
|
|
|
def test_prompt_is_a_first_class_resource_type():
|
|
"""``prompt`` is promoted out from under ``registered_model`` so RBAC can
|
|
address prompts directly. A grant on registered_model named ``foo`` must
|
|
not satisfy a request for prompt ``foo`` — that isolation lives in the
|
|
resolver, but the resource_type itself must exist in the valid set.
|
|
"""
|
|
assert "prompt" in VALID_RESOURCE_TYPES
|
|
assert "prompt" != "registered_model"
|
|
|
|
|
|
# ---- Scope-aware validator --------------------------------------------------
|
|
|
|
|
|
def test_grantable_permission_sets_pin_simplified_model():
|
|
"""Pin the two-tier workspace model and the resource-grant set.
|
|
|
|
Workspace grants accept only USE / MANAGE; resource grants accept any of
|
|
READ / USE / EDIT / MANAGE (no NO_PERMISSIONS). These sets gate every
|
|
permission write through the store, so a regression here would silently
|
|
re-enable the old five-value-anywhere model.
|
|
"""
|
|
assert WORKSPACE_GRANTABLE_PERMISSIONS == {USE.name, MANAGE.name}
|
|
assert RESOURCE_GRANTABLE_PERMISSIONS == {READ.name, USE.name, EDIT.name, MANAGE.name}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"permission",
|
|
[p.name for p in (READ, USE, EDIT, MANAGE)],
|
|
)
|
|
def test_validate_resource_grant_accepts_grantable(permission):
|
|
_validate_permission_for_resource_type(permission, "experiment")
|
|
_validate_permission_for_resource_type(permission, "registered_model")
|
|
_validate_permission_for_resource_type(permission, "prompt")
|
|
_validate_permission_for_resource_type(permission, "scorer")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"resource_type",
|
|
["experiment", "registered_model", "prompt", "scorer", "gateway_endpoint"],
|
|
)
|
|
def test_validate_resource_grant_rejects_no_permissions(resource_type):
|
|
with pytest.raises(MlflowException, match="NO_PERMISSIONS"):
|
|
_validate_permission_for_resource_type(NO_PERMISSIONS.name, resource_type)
|
|
|
|
|
|
@pytest.mark.parametrize("permission", [USE.name, MANAGE.name])
|
|
def test_validate_workspace_grant_accepts_use_or_manage(permission):
|
|
"""The unified workspace slot accepts USE (regular member) and MANAGE
|
|
(admin); the permission tier distinguishes the two without needing a
|
|
separate ``resource_type`` discriminant.
|
|
"""
|
|
_validate_permission_for_resource_type(permission, "workspace")
|
|
|
|
|
|
@pytest.mark.parametrize("permission", [READ.name, EDIT.name, NO_PERMISSIONS.name])
|
|
def test_validate_workspace_grant_rejects_other_tiers(permission):
|
|
with pytest.raises(MlflowException, match="resource_type='workspace'"):
|
|
_validate_permission_for_resource_type(permission, "workspace")
|
|
|
|
|
|
def test_validate_resource_type_rejects_legacy_any_discriminant():
|
|
"""The legacy ``resource_type='*'`` slot was retired in favor of the
|
|
unified ``'workspace'`` slot. Any write attempting the old shape should be
|
|
rejected at the validator level so stale callers fail loudly.
|
|
"""
|
|
with pytest.raises(MlflowException, match="Invalid resource type"):
|
|
_validate_permission_for_resource_type(USE.name, "*")
|
|
|
|
|
|
def test_validate_permission_for_resource_type_rejects_unknown():
|
|
with pytest.raises(MlflowException, match="Invalid permission"):
|
|
_validate_permission_for_resource_type("ADMIN", "experiment")
|
|
with pytest.raises(MlflowException, match="Invalid resource type"):
|
|
_validate_permission_for_resource_type(USE.name, "trace")
|