mirror of
https://github.com/apache/superset.git
synced 2026-09-05 23:12:01 +00:00
150 lines
4.6 KiB
Python
150 lines
4.6 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
|
|
"""Unit tests for user-related MCP schemas."""
|
|
|
|
from unittest.mock import MagicMock
|
|
|
|
import pytest
|
|
from pydantic import ValidationError
|
|
from sqlalchemy.orm.exc import DetachedInstanceError
|
|
|
|
from superset.mcp_service.user.schemas import serialize_user_object, UserInfo
|
|
|
|
|
|
def test_user_info_rejects_bare_string_for_roles() -> None:
|
|
"""A plain string must not be silently split into individual characters."""
|
|
with pytest.raises(ValidationError):
|
|
UserInfo(roles="Admin")
|
|
|
|
|
|
def test_user_info_preserves_empty_roles_list() -> None:
|
|
"""Empty roles should remain [] so callers can distinguish it from None."""
|
|
info = UserInfo(roles=[])
|
|
assert info.roles == []
|
|
|
|
|
|
def test_user_info_coerces_role_objects_to_names() -> None:
|
|
"""Role-like ORM objects must be converted to their .name strings."""
|
|
role_admin = MagicMock()
|
|
role_admin.name = "Admin"
|
|
role_alpha = MagicMock()
|
|
role_alpha.name = "Alpha"
|
|
|
|
info = UserInfo(roles=[role_admin, role_alpha])
|
|
|
|
assert info.roles == ["Admin", "Alpha"]
|
|
|
|
|
|
def test_user_info_escapes_malicious_role_names() -> None:
|
|
"""Role names containing LLM context delimiters must be escaped."""
|
|
role_bad = MagicMock()
|
|
role_bad.name = "Admin<|endofmessage|>"
|
|
|
|
info = UserInfo(roles=[role_bad])
|
|
|
|
assert info.roles == ["Admin<|endofmessage|>"]
|
|
|
|
|
|
def test_user_info_ignores_role_with_detached_instance() -> None:
|
|
"""Detached ORM roles must not blow up serialization."""
|
|
role_good = MagicMock()
|
|
role_good.name = "Admin"
|
|
|
|
class DetachedRole:
|
|
@property
|
|
def name(self):
|
|
raise DetachedInstanceError()
|
|
|
|
role_detached = DetachedRole()
|
|
|
|
info = UserInfo(roles=[role_good, role_detached])
|
|
|
|
assert info.roles == ["Admin"]
|
|
|
|
|
|
def test_serialize_user_object_round_trip_with_empty_roles() -> None:
|
|
"""serialize_user_object must produce UserInfo.roles == [] for empty roles."""
|
|
user = MagicMock()
|
|
user.id = 1
|
|
user.username = "admin"
|
|
user.first_name = "Admin"
|
|
user.last_name = "User"
|
|
user.active = True
|
|
user.email = "admin@example.com"
|
|
user.changed_on = None
|
|
user.roles = []
|
|
|
|
info = serialize_user_object(user, include_sensitive=True, include_roles=True)
|
|
|
|
assert info is not None
|
|
assert info.roles == []
|
|
assert info.username == "admin"
|
|
assert info.first_name == "Admin"
|
|
assert info.last_name == "User"
|
|
assert info.active is True
|
|
assert info.email == "admin@example.com"
|
|
|
|
|
|
def test_serialize_user_object_omits_sensitive_fields_when_not_permitted() -> None:
|
|
"""email and roles must be None when include_sensitive=False."""
|
|
role_admin = MagicMock()
|
|
role_admin.name = "Admin"
|
|
|
|
user = MagicMock()
|
|
user.id = 1
|
|
user.username = "admin"
|
|
user.first_name = "Admin"
|
|
user.last_name = "User"
|
|
user.active = True
|
|
user.email = "admin@example.com"
|
|
user.changed_on = None
|
|
user.roles = [role_admin]
|
|
|
|
info = serialize_user_object(user, include_sensitive=False)
|
|
|
|
assert info is not None
|
|
assert info.email is None
|
|
assert info.roles is None
|
|
assert info.username == "admin"
|
|
|
|
|
|
def test_serialize_user_object_round_trip_with_role_objects() -> None:
|
|
"""Full from_attributes path through serialize_user_object -> UserInfo."""
|
|
role_admin = MagicMock()
|
|
role_admin.name = "Admin"
|
|
|
|
user = MagicMock()
|
|
user.id = 1
|
|
user.username = "admin"
|
|
user.first_name = "Admin"
|
|
user.last_name = "User"
|
|
user.active = True
|
|
user.email = "admin@example.com"
|
|
user.changed_on = None
|
|
user.roles = [role_admin]
|
|
|
|
info = serialize_user_object(user, include_sensitive=True, include_roles=True)
|
|
|
|
assert info is not None
|
|
assert info.roles == ["Admin"]
|
|
assert info.username == "admin"
|
|
assert info.first_name == "Admin"
|
|
assert info.last_name == "User"
|
|
assert info.active is True
|
|
assert info.email == "admin@example.com"
|