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

# Unit tests for the review-queue / label-schema authorization validators: they
# exercise the gate decisions in ``mlflow.server.auth`` directly by mocking the
# request / store / experiment-permission boundary, so no live auth server is
# needed. End-to-end routing is covered by the broader auth client tests.
import json
from types import SimpleNamespace

import pytest

from mlflow.exceptions import MlflowException
from mlflow.genai.review_queues import ReviewQueueType
from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE, ErrorCode
from mlflow.protos.review_queues_pb2 import ListReviewQueues
from mlflow.utils.proto_json_utils import message_to_json, parse_dict

# flask-wtf (the auth extra) is required to import the auth app.
pytest.importorskip("flask_wtf")

from mlflow.server import auth
from mlflow.server.auth.permissions import get_permission


def _setup(
    monkeypatch,
    *,
    permission,
    queue_users=None,
    username="alice",
    created_by=None,
    queue_type=ReviewQueueType.CUSTOM,
    update_body=None,
):
    """Patch the request / store / permission boundary the validators read."""
    perm = get_permission(permission)
    queue = SimpleNamespace(
        experiment_id="123",
        users=list(queue_users or []),
        created_by=created_by,
        queue_type=queue_type,
    )
    schema = SimpleNamespace(experiment_id="123")
    store = SimpleNamespace(
        get_review_queue=lambda _qid: queue,
        get_review_queue_by_name=lambda _exp, name: queue,
        get_label_schema=lambda _sid: schema,
    )
    monkeypatch.setattr(auth, "authenticate_request", lambda: SimpleNamespace(username=username))
    monkeypatch.setattr(
        auth,
        "_get_request_param",
        lambda p: {"queue_id": "q1", "schema_id": "s1", "experiment_id": "123", "name": "Q"}[p],
    )
    monkeypatch.setattr(auth, "_get_experiment_permission", lambda _exp, _user: perm)
    monkeypatch.setattr(auth, "_get_tracking_store", lambda: store)
    # The owner-reassignment gate parses the live request body; stub it so the
    # real `_update_review_queue_reassigns_owner` detection runs against it.
    monkeypatch.setattr(
        auth, "request", SimpleNamespace(get_json=lambda silent=False: dict(update_body or {}))
    )


def test_review_queue_has_member_normalizes_incoming_username():
    # Assigned users are normalized (lowercased + stripped) at write time, so the
    # helper normalizes only the incoming username and matches it directly.
    queue = SimpleNamespace(users=["alice", "bob"])
    assert auth._review_queue_has_member(queue, "Alice")
    assert auth._review_queue_has_member(queue, " BOB ")
    assert not auth._review_queue_has_member(queue, "carol")
    assert not auth._review_queue_has_member(SimpleNamespace(users=[]), "alice")


def test_is_review_queue_owner_is_case_insensitive():
    assert auth._is_review_queue_owner(SimpleNamespace(created_by="Alice"), "alice")
    assert auth._is_review_queue_owner(SimpleNamespace(created_by=" alice "), "ALICE")
    assert not auth._is_review_queue_owner(SimpleNamespace(created_by="bob"), "alice")
    # No owner (no-auth-era rows) is never an owner match.
    assert not auth._is_review_queue_owner(SimpleNamespace(created_by=None), "alice")
    assert not auth._is_review_queue_owner(SimpleNamespace(created_by=""), "alice")


@pytest.mark.parametrize(
    "validator",
    ["validate_can_create_label_schema", "validate_can_manage_label_schema"],
)
@pytest.mark.parametrize(
    ("permission", "expected"), [("READ", False), ("EDIT", False), ("MANAGE", True)]
)
def test_label_schema_management_requires_manage(monkeypatch, validator, permission, expected):
    _setup(monkeypatch, permission=permission)
    assert getattr(auth, validator)() is expected


@pytest.mark.parametrize(
    ("permission", "expected"), [("READ", False), ("EDIT", True), ("MANAGE", True)]
)
def test_create_add_items_and_user_queue_require_edit(monkeypatch, permission, expected):
    # Creating (and owning) a queue, flagging items, and routing to a user queue
    # are all experiment-EDIT operations now.
    _setup(monkeypatch, permission=permission)
    assert auth.validate_can_create_review_queue() is expected
    assert auth.validate_can_add_items_to_review_queue() is expected
    assert auth.validate_can_get_or_create_user_queue() is expected


def test_update_and_remove_items_allow_owner_or_manager(monkeypatch):
    # Manager edits any queue.
    _setup(monkeypatch, permission="MANAGE", created_by="bob", username="alice")
    assert auth.validate_can_update_review_queue() is True
    assert auth.validate_can_remove_items_from_review_queue() is True

    # Owning EDIT user edits their own queue.
    _setup(monkeypatch, permission="EDIT", created_by="alice", username="alice")
    assert auth.validate_can_update_review_queue() is True
    assert auth.validate_can_remove_items_from_review_queue() is True

    # EDIT non-owner is denied.
    _setup(monkeypatch, permission="EDIT", created_by="bob", username="alice")
    assert auth.validate_can_update_review_queue() is False
    assert auth.validate_can_remove_items_from_review_queue() is False

    # READ owner is denied (ownership amplifies EDIT, never substitutes).
    _setup(monkeypatch, permission="READ", created_by="alice", username="alice")
    assert auth.validate_can_update_review_queue() is False


def test_remove_items_on_user_queue_is_manage_only(monkeypatch):
    # A manager may prune (remove items from) a personal USER queue...
    _setup(
        monkeypatch,
        permission="MANAGE",
        created_by="alice",
        username="alice",
        queue_type=ReviewQueueType.USER,
    )
    assert auth.validate_can_remove_items_from_review_queue() is True

    # ...but an EDIT user can't un-assign work from their own USER queue (its
    # contents are a manager's call, matching the USER-queue delete rule).
    _setup(
        monkeypatch,
        permission="EDIT",
        created_by="alice",
        username="alice",
        queue_type=ReviewQueueType.USER,
    )
    assert auth.validate_can_remove_items_from_review_queue() is False


def test_owner_reassignment_requires_manage(monkeypatch):
    # An owning EDIT user may edit shape but NOT reassign the owner.
    _setup(
        monkeypatch,
        permission="EDIT",
        created_by="alice",
        username="alice",
        update_body={"queue_id": "q1", "new_owner": "victim"},
    )
    assert auth.validate_can_update_review_queue() is False

    # A manager may reassign the owner.
    _setup(
        monkeypatch,
        permission="MANAGE",
        created_by="bob",
        username="alice",
        update_body={"queue_id": "q1", "new_owner": "victim"},
    )
    assert auth.validate_can_update_review_queue() is True


def _users(*names):
    return [SimpleNamespace(username=n) for n in names]


def _forbid_list_users():
    raise AssertionError("store.list_users must not be called here")


def test_create_custom_queue_rejects_registered_username(monkeypatch):
    # A custom queue named after a registered user would shadow that user's
    # personal queue; reject it case-insensitively ("Alice" vs user "alice").
    _setup(
        monkeypatch,
        permission="EDIT",
        update_body={"experiment_id": "123", "name": "Alice", "queue_type": "CUSTOM"},
    )
    monkeypatch.setattr(auth.store, "list_users", lambda: _users("alice", "bob"))
    with pytest.raises(MlflowException, match="registered user") as exc:
        auth.validate_can_create_review_queue()
    assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)


def test_create_custom_queue_allows_non_username(monkeypatch):
    _setup(
        monkeypatch,
        permission="EDIT",
        update_body={"experiment_id": "123", "name": "Hallucinations", "queue_type": "CUSTOM"},
    )
    consulted = []
    monkeypatch.setattr(
        auth.store, "list_users", lambda: consulted.append(True) or _users("alice", "bob")
    )
    assert auth.validate_can_create_review_queue() is True
    # The roster was actually consulted (guards against the match silently no-opping).
    assert consulted == [True]


def test_create_unset_queue_type_skips_shadow_check(monkeypatch):
    # No queue_type (proto 0/UNSPECIFIED) skips the check without calling list_users
    # and without `from_proto` raising on the unspecified value.
    _setup(monkeypatch, permission="EDIT", update_body={"experiment_id": "123", "name": "alice"})
    monkeypatch.setattr(auth.store, "list_users", _forbid_list_users)
    assert auth.validate_can_create_review_queue() is True


def test_create_blank_name_skips_shadow_check(monkeypatch):
    # A blank name can't match a user; skip before querying the roster (the
    # handler/store validate the empty name separately).
    _setup(
        monkeypatch,
        permission="EDIT",
        update_body={"experiment_id": "123", "name": "   ", "queue_type": "CUSTOM"},
    )
    monkeypatch.setattr(auth.store, "list_users", _forbid_list_users)
    assert auth.validate_can_create_review_queue() is True


def test_create_user_queue_named_after_user_is_allowed(monkeypatch):
    # A USER queue *is* its username, so the CUSTOM-only shadowing guard skips it.
    _setup(
        monkeypatch,
        permission="EDIT",
        update_body={"experiment_id": "123", "name": "alice", "queue_type": "USER"},
    )
    monkeypatch.setattr(auth.store, "list_users", _forbid_list_users)
    assert auth.validate_can_create_review_queue() is True


def test_create_shadow_check_skipped_when_unauthorized(monkeypatch):
    # The permission gate denies first; the name check (and user-list query) never runs.
    _setup(
        monkeypatch,
        permission="READ",
        update_body={"experiment_id": "123", "name": "alice", "queue_type": "CUSTOM"},
    )
    monkeypatch.setattr(auth.store, "list_users", _forbid_list_users)
    assert auth.validate_can_create_review_queue() is False


def test_rename_custom_queue_to_username_rejected(monkeypatch):
    _setup(
        monkeypatch,
        permission="MANAGE",
        created_by="bob",
        username="alice",
        update_body={"queue_id": "q1", "name": "Bob"},
    )
    monkeypatch.setattr(auth.store, "list_users", lambda: _users("alice", "bob"))
    with pytest.raises(MlflowException, match="registered user") as exc:
        auth.validate_can_update_review_queue()
    assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)


def test_rename_custom_queue_to_non_username_allowed(monkeypatch):
    _setup(
        monkeypatch,
        permission="MANAGE",
        created_by="bob",
        username="alice",
        update_body={"queue_id": "q1", "name": "Renamed"},
    )
    monkeypatch.setattr(auth.store, "list_users", lambda: _users("alice", "bob"))
    assert auth.validate_can_update_review_queue() is True


def test_rename_shadow_guard_skips_user_queues(monkeypatch):
    # Renaming a USER queue is rejected by the store, not this guard; the guard is
    # CUSTOM-only, so it neither fires nor queries the user list here.
    _setup(
        monkeypatch,
        permission="MANAGE",
        created_by="alice",
        username="alice",
        queue_type=ReviewQueueType.USER,
        update_body={"queue_id": "q1", "name": "bob"},
    )
    monkeypatch.setattr(auth.store, "list_users", _forbid_list_users)
    assert auth.validate_can_update_review_queue() is True


def _route_for(validator):
    # The exact (path, method) the dispatcher would match to this validator.
    return next(
        (path, method)
        for (path, method), v in auth.BEFORE_REQUEST_VALIDATORS.items()
        if v is validator
    )


def _patch_request(monkeypatch, validator, body):
    path, method = _route_for(validator)
    monkeypatch.setattr(
        auth,
        "request",
        SimpleNamespace(path=path, method=method, get_json=lambda silent=False: dict(body)),
    )


def test_admin_create_shadowing_username_is_blocked(monkeypatch):
    # Admins bypass the validators, so `_before_request` runs the integrity guard
    # directly for them; shadowing a username is still rejected. permission is
    # irrelevant here (admins never hit the permission gate).
    body = {"experiment_id": "123", "name": "Bob", "queue_type": "CUSTOM"}
    _setup(monkeypatch, permission="READ", update_body=body)
    _patch_request(monkeypatch, auth.validate_can_create_review_queue, body)
    monkeypatch.setattr(auth.store, "list_users", lambda: _users("alice", "bob"))
    with pytest.raises(MlflowException, match="registered user") as exc:
        auth.enforce_review_queue_name_not_username()
    assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)


def test_admin_rename_onto_username_is_blocked(monkeypatch):
    body = {"queue_id": "q1", "name": "Bob"}
    _setup(monkeypatch, permission="READ", created_by="bob", update_body=body)
    _patch_request(monkeypatch, auth.validate_can_update_review_queue, body)
    monkeypatch.setattr(auth.store, "list_users", lambda: _users("alice", "bob"))
    with pytest.raises(MlflowException, match="registered user") as exc:
        auth.enforce_review_queue_name_not_username()
    assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE)


def test_admin_create_non_shadowing_name_is_allowed(monkeypatch):
    body = {"experiment_id": "123", "name": "Hallucinations", "queue_type": "CUSTOM"}
    _setup(monkeypatch, permission="READ", update_body=body)
    _patch_request(monkeypatch, auth.validate_can_create_review_queue, body)
    monkeypatch.setattr(auth.store, "list_users", lambda: _users("alice", "bob"))
    auth.enforce_review_queue_name_not_username()


def test_name_guard_is_a_noop_off_the_review_queue_routes(monkeypatch):
    # On any other route the guard does nothing and never touches the roster, so
    # adding it to `_before_request`'s admin branch is free for every endpoint.
    _setup(monkeypatch, permission="READ")
    _patch_request(monkeypatch, auth.validate_can_read_experiment, {})
    monkeypatch.setattr(auth.store, "list_users", _forbid_list_users)
    auth.enforce_review_queue_name_not_username()


def test_admin_owner_reassignment_without_rename_is_a_noop(monkeypatch):
    # An admin update that only reassigns the owner (no `name`) is not a rename, so
    # the guard short-circuits before querying the roster.
    body = {"queue_id": "q1", "new_owner": "victim"}
    _setup(monkeypatch, permission="READ", created_by="bob", update_body=body)
    _patch_request(monkeypatch, auth.validate_can_update_review_queue, body)
    monkeypatch.setattr(auth.store, "list_users", _forbid_list_users)
    auth.enforce_review_queue_name_not_username()


def test_admin_owner_reassignment_with_shadowing_rename_is_blocked(monkeypatch):
    # Reassigning the owner and renaming onto a username in the same request: the
    # rename still shadows, so it's rejected even alongside `new_owner`.
    body = {"queue_id": "q1", "new_owner": "victim", "name": "Bob"}
    _setup(monkeypatch, permission="READ", created_by="bob", update_body=body)
    _patch_request(monkeypatch, auth.validate_can_update_review_queue, body)
    monkeypatch.setattr(auth.store, "list_users", lambda: _users("alice", "bob"))
    with pytest.raises(MlflowException, match="registered user"):
        auth.enforce_review_queue_name_not_username()


def test_admin_create_shadowing_username_via_camelcase_is_blocked(monkeypatch):
    # Protobuf JSON accepts camelCase keys; the guard parses the proto (not raw
    # keys), so `queueType` is detected the same way the handler reads it.
    body = {"experimentId": "123", "name": "Bob", "queueType": "CUSTOM"}
    _setup(monkeypatch, permission="READ", update_body=body)
    _patch_request(monkeypatch, auth.validate_can_create_review_queue, body)
    monkeypatch.setattr(auth.store, "list_users", lambda: _users("alice", "bob"))
    with pytest.raises(MlflowException, match="registered user"):
        auth.enforce_review_queue_name_not_username()


def test_owner_reassignment_via_camelcase_still_requires_manage(monkeypatch):
    # Regression: protobuf JSON also accepts the camelCase `newOwner`. The gate
    # must detect it by parsing the proto (not scanning raw JSON keys), so an
    # owning EDIT user can't dodge the MANAGE-only owner reassignment.
    _setup(
        monkeypatch,
        permission="EDIT",
        created_by="alice",
        username="alice",
        update_body={"queue_id": "q1", "newOwner": "victim"},
    )
    assert auth.validate_can_update_review_queue() is False


def test_delete_owner_can_delete_own_custom_but_not_user(monkeypatch):
    # Manager deletes any queue.
    _setup(monkeypatch, permission="MANAGE", created_by="bob", username="alice")
    assert auth.validate_can_delete_review_queue() is True

    # Owning EDIT user deletes their own CUSTOM queue.
    _setup(monkeypatch, permission="EDIT", created_by="alice", username="alice")
    assert auth.validate_can_delete_review_queue() is True

    # ...but not a USER queue (those are MANAGE-only to delete).
    _setup(
        monkeypatch,
        permission="EDIT",
        created_by="alice",
        username="alice",
        queue_type=ReviewQueueType.USER,
    )
    assert auth.validate_can_delete_review_queue() is False

    # EDIT non-owner is denied.
    _setup(monkeypatch, permission="EDIT", created_by="bob", username="alice")
    assert auth.validate_can_delete_review_queue() is False


@pytest.mark.parametrize(
    ("permission", "expected"), [("READ", True), ("EDIT", True), ("MANAGE", True)]
)
def test_read_label_schema_requires_read(monkeypatch, permission, expected):
    _setup(monkeypatch, permission=permission)
    assert auth.validate_can_read_label_schema() is expected


def test_review_queue_item_requires_edit_and_membership(monkeypatch):
    # EDIT + assigned -> allowed.
    _setup(monkeypatch, permission="EDIT", queue_users=["alice"], username="alice")
    assert auth.validate_can_review_queue_item() is True

    # EDIT but not assigned -> denied (even a manager must self-assign first).
    _setup(monkeypatch, permission="MANAGE", queue_users=["bob"], username="alice")
    assert auth.validate_can_review_queue_item() is False

    # Assigned but only READ -> denied (reviewing needs EDIT).
    _setup(monkeypatch, permission="READ", queue_users=["alice"], username="alice")
    assert auth.validate_can_review_queue_item() is False


def test_view_review_queue_manage_owner_or_membership(monkeypatch):
    # Manager sees any queue, assigned or not.
    _setup(monkeypatch, permission="MANAGE", queue_users=["bob"], username="alice")
    assert auth.validate_can_view_review_queue() is True

    # READ + assigned -> visible.
    _setup(monkeypatch, permission="READ", queue_users=["alice"], username="alice")
    assert auth.validate_can_view_review_queue() is True

    # READ + not assigned -> hidden.
    _setup(monkeypatch, permission="READ", queue_users=["bob"], username="alice")
    assert auth.validate_can_view_review_queue() is False

    # EDIT owner (not assigned) -> visible.
    _setup(
        monkeypatch, permission="EDIT", queue_users=["bob"], created_by="alice", username="alice"
    )
    assert auth.validate_can_view_review_queue() is True

    # READ owner -> hidden (ownership amplifies EDIT, never substitutes).
    _setup(
        monkeypatch, permission="READ", queue_users=["bob"], created_by="alice", username="alice"
    )
    assert auth.validate_can_view_review_queue() is False


def test_view_review_queue_by_name_manage_or_membership(monkeypatch):
    # Manager sees any queue by name.
    _setup(monkeypatch, permission="MANAGE", queue_users=["bob"], username="alice")
    assert auth.validate_can_view_review_queue_by_name() is True

    # READ + assigned -> visible; READ + not assigned -> hidden.
    _setup(monkeypatch, permission="READ", queue_users=["alice"], username="alice")
    assert auth.validate_can_view_review_queue_by_name() is True
    _setup(monkeypatch, permission="READ", queue_users=["bob"], username="alice")
    assert auth.validate_can_view_review_queue_by_name() is False

    # EDIT owner -> visible; READ owner -> hidden.
    _setup(
        monkeypatch, permission="EDIT", queue_users=["bob"], created_by="alice", username="alice"
    )
    assert auth.validate_can_view_review_queue_by_name() is True
    _setup(
        monkeypatch, permission="READ", queue_users=["bob"], created_by="alice", username="alice"
    )
    assert auth.validate_can_view_review_queue_by_name() is False


def _filter_response(monkeypatch, *, permission, username, queues):
    _setup(monkeypatch, permission=permission, username=username)
    monkeypatch.setattr(auth, "sender_is_admin", lambda: False)
    msg = ListReviewQueues.Response()
    parse_dict({"review_queues": queues}, msg)
    resp = SimpleNamespace(json=json.loads(message_to_json(msg)), data=None)
    auth.filter_list_review_queues(resp)
    return resp


def test_filter_list_review_queues_read_only_sees_only_assigned(monkeypatch):
    resp = _filter_response(
        monkeypatch,
        permission="READ",
        username="alice",
        queues=[{"queue_id": "q1", "users": ["alice"]}, {"queue_id": "q2", "users": ["bob"]}],
    )
    out = ListReviewQueues.Response()
    parse_dict(json.loads(resp.data), out)
    assert [q.queue_id for q in out.review_queues] == ["q1"]


@pytest.mark.parametrize("permission", ["EDIT", "MANAGE"])
def test_filter_list_review_queues_editor_and_manager_see_all(monkeypatch, permission):
    resp = _filter_response(
        monkeypatch,
        permission=permission,
        username="alice",
        queues=[{"queue_id": "q1", "users": ["bob"]}, {"queue_id": "q2", "users": ["carol"]}],
    )
    # EDIT and MANAGE both short-circuit before rewriting the response (all rows).
    assert resp.data is None