mirror of
https://github.com/langgenius/dify.git
synced 2026-09-03 06:53:26 +08:00
test: migrate controller auth wrapper sessions and ORM models to SQLite (#40529)
This commit is contained in:
Asuka Minato
GitHub
parent
2aaac4dadb
commit
a575ae5ef2
@ -1,8 +1,9 @@
|
||||
from unittest.mock import MagicMock, patch
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
from sqlalchemy.orm import Session
|
||||
from sqlalchemy.orm import Session, scoped_session
|
||||
from werkzeug.exceptions import Forbidden, NotFound
|
||||
|
||||
import controllers.console.explore.wraps as wraps_module
|
||||
@ -20,14 +21,20 @@ from controllers.console.explore.wraps import (
|
||||
trial_feature_enable,
|
||||
user_allowed_to_access_app,
|
||||
)
|
||||
from models import AccountTrialAppRecord, App, AppMode, InstalledApp, TrialApp
|
||||
from models import Account, AccountTrialAppRecord, App, AppMode, InstalledApp, TrialApp
|
||||
|
||||
|
||||
def _bind_database(monkeypatch: pytest.MonkeyPatch, sqlite_session: Session) -> None:
|
||||
session_proxy = MagicMock(wraps=sqlite_session)
|
||||
session_proxy.return_value = sqlite_session
|
||||
monkeypatch.setattr(wraps_module.db, "session", session_proxy)
|
||||
monkeypatch.setattr(model_module.db, "session", session_proxy)
|
||||
session_registry = scoped_session(lambda: sqlite_session)
|
||||
monkeypatch.setattr(wraps_module.db, "session", session_registry)
|
||||
monkeypatch.setattr(model_module.db, "session", session_registry)
|
||||
|
||||
|
||||
def _account(*, account_id: str | None = None) -> Account:
|
||||
account = Account(name="Explore user", email="user@example.com")
|
||||
if account_id is not None:
|
||||
account.id = account_id
|
||||
return account
|
||||
|
||||
|
||||
def _app() -> App:
|
||||
@ -67,7 +74,7 @@ def test_installed_app_required_not_found(
|
||||
|
||||
with patch(
|
||||
"controllers.console.explore.wraps.current_account_with_tenant",
|
||||
return_value=(MagicMock(), tenant_id),
|
||||
return_value=(_account(), tenant_id),
|
||||
):
|
||||
with pytest.raises(NotFound):
|
||||
view(str(uuid4()))
|
||||
@ -91,7 +98,7 @@ def test_installed_app_required_app_deleted(
|
||||
|
||||
with patch(
|
||||
"controllers.console.explore.wraps.current_account_with_tenant",
|
||||
return_value=(MagicMock(), tenant_id),
|
||||
return_value=(_account(), tenant_id),
|
||||
):
|
||||
with pytest.raises(NotFound):
|
||||
view(installed_app_id)
|
||||
@ -116,7 +123,7 @@ def test_installed_app_required_success(
|
||||
|
||||
with patch(
|
||||
"controllers.console.explore.wraps.current_account_with_tenant",
|
||||
return_value=(MagicMock(), app.tenant_id),
|
||||
return_value=(_account(), app.tenant_id),
|
||||
):
|
||||
result = view(installed_app.id)
|
||||
|
||||
@ -126,19 +133,18 @@ def test_installed_app_required_success(
|
||||
|
||||
|
||||
def test_user_allowed_to_access_app_denied():
|
||||
installed_app = MagicMock(app_id="app-1")
|
||||
installed_app = _installed_app(app_id="app-1", tenant_id="tenant-1")
|
||||
|
||||
@user_allowed_to_access_app
|
||||
def view(installed_app):
|
||||
return "ok"
|
||||
|
||||
feature = MagicMock()
|
||||
feature.webapp_auth.enabled = True
|
||||
feature = SimpleNamespace(webapp_auth=SimpleNamespace(enabled=True))
|
||||
|
||||
with (
|
||||
patch(
|
||||
"controllers.console.explore.wraps.current_account_with_tenant",
|
||||
return_value=(MagicMock(id="user-1"), None),
|
||||
return_value=(_account(account_id="user-1"), None),
|
||||
),
|
||||
patch(
|
||||
"controllers.console.explore.wraps.FeatureService.get_system_features",
|
||||
@ -154,19 +160,18 @@ def test_user_allowed_to_access_app_denied():
|
||||
|
||||
|
||||
def test_user_allowed_to_access_app_success():
|
||||
installed_app = MagicMock(app_id="app-1")
|
||||
installed_app = _installed_app(app_id="app-1", tenant_id="tenant-1")
|
||||
|
||||
@user_allowed_to_access_app
|
||||
def view(installed_app):
|
||||
return "ok"
|
||||
|
||||
feature = MagicMock()
|
||||
feature.webapp_auth.enabled = True
|
||||
feature = SimpleNamespace(webapp_auth=SimpleNamespace(enabled=True))
|
||||
|
||||
with (
|
||||
patch(
|
||||
"controllers.console.explore.wraps.current_account_with_tenant",
|
||||
return_value=(MagicMock(id="user-1"), None),
|
||||
return_value=(_account(account_id="user-1"), None),
|
||||
),
|
||||
patch(
|
||||
"controllers.console.explore.wraps.FeatureService.get_system_features",
|
||||
@ -193,7 +198,7 @@ def test_trial_app_required_not_allowed(
|
||||
|
||||
with patch(
|
||||
"controllers.console.explore.wraps.current_account_with_tenant",
|
||||
return_value=(MagicMock(id=str(uuid4())), None),
|
||||
return_value=(_account(account_id=str(uuid4())), None),
|
||||
):
|
||||
with pytest.raises(TrialAppNotAllowed):
|
||||
view(str(uuid4()))
|
||||
@ -218,7 +223,7 @@ def test_trial_app_required_limit_exceeded(
|
||||
|
||||
with patch(
|
||||
"controllers.console.explore.wraps.current_account_with_tenant",
|
||||
return_value=(MagicMock(id=account_id), None),
|
||||
return_value=(_account(account_id=account_id), None),
|
||||
):
|
||||
with pytest.raises(TrialAppLimitExceeded):
|
||||
view(app.id)
|
||||
@ -243,7 +248,7 @@ def test_trial_app_required_success(
|
||||
|
||||
with patch(
|
||||
"controllers.console.explore.wraps.current_account_with_tenant",
|
||||
return_value=(MagicMock(id=account_id), None),
|
||||
return_value=(_account(account_id=account_id), None),
|
||||
):
|
||||
result = view(app.id)
|
||||
|
||||
|
||||
@ -7,7 +7,8 @@ from uuid import NAMESPACE_URL, uuid5
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from sqlalchemy.orm import Session, sessionmaker
|
||||
from sqlalchemy import Engine, event
|
||||
from sqlalchemy.orm import Session
|
||||
from werkzeug.exceptions import HTTPException
|
||||
|
||||
from configs import dify_config
|
||||
@ -264,7 +265,7 @@ class TestEnterpriseInnerApiUserAuth:
|
||||
# Assert
|
||||
assert result == "no_user"
|
||||
|
||||
def test_should_pass_through_when_hmac_signature_invalid(self, app: Flask):
|
||||
def test_should_pass_through_when_hmac_signature_invalid(self, app: Flask, sqlite_engine: Engine):
|
||||
"""Invalid HMAC auth passes through without opening a database session."""
|
||||
|
||||
# Arrange
|
||||
@ -272,17 +273,20 @@ class TestEnterpriseInnerApiUserAuth:
|
||||
def protected_view(**kwargs):
|
||||
return kwargs.get("user", "no_user")
|
||||
|
||||
# Act - use wrong signature
|
||||
with app.test_request_context(
|
||||
headers={"Authorization": "Bearer user123:wrong_signature", "X-Inner-Api-Key": "valid_key"}
|
||||
):
|
||||
with patch.object(dify_config, "INNER_API", True):
|
||||
with patch("controllers.inner_api.wraps.session_factory.create_session") as mock_create_session:
|
||||
result = protected_view()
|
||||
def fail_on_query(*_args, **_kwargs):
|
||||
pytest.fail("invalid HMAC must not access the database")
|
||||
|
||||
event.listen(sqlite_engine, "before_cursor_execute", fail_on_query)
|
||||
try:
|
||||
with app.test_request_context(
|
||||
headers={"Authorization": "Bearer user123:wrong_signature", "X-Inner-Api-Key": "valid_key"}
|
||||
):
|
||||
with patch.object(dify_config, "INNER_API", True):
|
||||
result = protected_view()
|
||||
finally:
|
||||
event.remove(sqlite_engine, "before_cursor_execute", fail_on_query)
|
||||
|
||||
# Assert
|
||||
assert result == "no_user"
|
||||
mock_create_session.assert_not_called()
|
||||
|
||||
@pytest.mark.parametrize("sqlite_session", [(EndUser,)], indirect=True)
|
||||
def test_should_inject_user_when_hmac_signature_valid(self, app: Flask, sqlite_session: Session):
|
||||
@ -312,21 +316,13 @@ class TestEnterpriseInnerApiUserAuth:
|
||||
)
|
||||
sqlite_session.add(end_user)
|
||||
sqlite_session.commit()
|
||||
database_session_factory = sessionmaker(
|
||||
bind=sqlite_session.get_bind(),
|
||||
expire_on_commit=False,
|
||||
)
|
||||
|
||||
# Act
|
||||
with app.test_request_context(
|
||||
headers={"Authorization": f"Bearer {user_id}:{valid_signature}", "X-Inner-Api-Key": inner_api_key}
|
||||
):
|
||||
with patch.object(dify_config, "INNER_API", True):
|
||||
with patch(
|
||||
"controllers.inner_api.wraps.session_factory.create_session",
|
||||
database_session_factory,
|
||||
):
|
||||
result = protected_view()
|
||||
result = protected_view()
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, EndUser)
|
||||
|
||||
@ -1,11 +1,12 @@
|
||||
import uuid
|
||||
from contextlib import nullcontext
|
||||
from unittest.mock import MagicMock, patch
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import PropertyMock, patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from sqlalchemy import Engine, select
|
||||
from sqlalchemy.orm import Session
|
||||
from werkzeug.exceptions import Forbidden, NotFound, Unauthorized
|
||||
from werkzeug.exceptions import Forbidden, InternalServerError, NotFound, Unauthorized
|
||||
|
||||
from controllers.openapi.auth.data import AuthData, ExternalIdentity
|
||||
from controllers.openapi.auth.prepare import (
|
||||
@ -18,346 +19,347 @@ from controllers.openapi.auth.prepare import (
|
||||
resolve_external_user,
|
||||
)
|
||||
from libs.oauth_bearer import TokenType
|
||||
from models.account import TenantAccountRole
|
||||
from models import Account, App, EndUser, Tenant, TenantAccountJoin
|
||||
from models.account import AccountStatus, TenantAccountRole, TenantStatus
|
||||
from models.enums import AppStatus
|
||||
from models.model import AppMode, IconType
|
||||
from services import end_user_service
|
||||
from services.enterprise.enterprise_service import WebAppAccessMode
|
||||
|
||||
APP_ID = "00000000-0000-0000-0000-000000000001"
|
||||
TENANT_ID = "00000000-0000-0000-0000-000000000002"
|
||||
ACCOUNT_ID = "00000000-0000-0000-0000-000000000003"
|
||||
|
||||
|
||||
def _make_auth_data(**kwargs) -> AuthData:
|
||||
mock_fields = {k: kwargs.pop(k) for k in ("app", "tenant", "caller") if k in kwargs}
|
||||
data = AuthData(
|
||||
def _make_auth_data(**kwargs: object) -> AuthData:
|
||||
return AuthData(
|
||||
token_type=kwargs.pop("token_type", TokenType.OAUTH_ACCOUNT),
|
||||
token_hash=kwargs.pop("token_hash", "testhash"),
|
||||
scopes=kwargs.pop("scopes", frozenset()),
|
||||
**kwargs,
|
||||
)
|
||||
for k, v in mock_fields.items():
|
||||
setattr(data, k, v)
|
||||
return data
|
||||
|
||||
|
||||
_VALID_APP_UUID = "00000000-0000-0000-0000-000000000001"
|
||||
|
||||
|
||||
def test_load_app_writes_app_to_data():
|
||||
app = MagicMock()
|
||||
app.status = "normal"
|
||||
app.enable_api = True
|
||||
data = _make_auth_data(path_params={"app_id": _VALID_APP_UUID})
|
||||
with patch("controllers.openapi.auth.prepare.AppService.get_app_by_id", return_value=app):
|
||||
load_app(data)
|
||||
assert data.app is app
|
||||
|
||||
|
||||
def test_load_app_raises_not_found_for_non_uuid_app_id():
|
||||
data = _make_auth_data(path_params={"app_id": "not-a-uuid"})
|
||||
with pytest.raises(NotFound):
|
||||
load_app(data)
|
||||
|
||||
|
||||
def test_load_app_raises_not_found_when_missing():
|
||||
data = _make_auth_data(path_params={"app_id": _VALID_APP_UUID})
|
||||
with patch("controllers.openapi.auth.prepare.AppService.get_app_by_id", return_value=None):
|
||||
with pytest.raises(NotFound):
|
||||
load_app(data)
|
||||
|
||||
|
||||
def test_load_app_raises_not_found_when_not_normal():
|
||||
app = MagicMock()
|
||||
app.status = "archived"
|
||||
data = _make_auth_data(path_params={"app_id": _VALID_APP_UUID})
|
||||
with patch("controllers.openapi.auth.prepare.AppService.get_app_by_id", return_value=app):
|
||||
with pytest.raises(NotFound):
|
||||
load_app(data)
|
||||
|
||||
|
||||
def test_load_app_stashes_app_even_when_api_disabled():
|
||||
app = MagicMock()
|
||||
app.status = "normal"
|
||||
app.enable_api = False
|
||||
data = _make_auth_data(path_params={"app_id": _VALID_APP_UUID})
|
||||
with patch("controllers.openapi.auth.prepare.AppService.get_app_by_id", return_value=app):
|
||||
load_app(data)
|
||||
assert data.app is app
|
||||
|
||||
|
||||
def test_load_app_skips_when_already_set():
|
||||
existing_app = MagicMock()
|
||||
data = _make_auth_data(app=existing_app, path_params={"app_id": "abc"})
|
||||
load_app(data)
|
||||
assert data.app is existing_app
|
||||
|
||||
|
||||
def test_load_tenant_writes_tenant():
|
||||
app = MagicMock()
|
||||
app.tenant_id = uuid.uuid4()
|
||||
tenant = MagicMock()
|
||||
tenant.status = "normal"
|
||||
data = _make_auth_data(app=app)
|
||||
with patch("controllers.openapi.auth.prepare.TenantService.get_tenant_by_id", return_value=tenant):
|
||||
load_tenant(data)
|
||||
assert data.tenant is tenant
|
||||
|
||||
|
||||
def test_load_tenant_skips_when_already_set():
|
||||
existing_tenant = MagicMock()
|
||||
data = _make_auth_data(app=MagicMock(), tenant=existing_tenant)
|
||||
load_tenant(data)
|
||||
assert data.tenant is existing_tenant
|
||||
|
||||
|
||||
def test_load_tenant_raises_forbidden_when_archived():
|
||||
from models.account import TenantStatus
|
||||
|
||||
app = MagicMock()
|
||||
app.tenant_id = uuid.uuid4()
|
||||
tenant = MagicMock()
|
||||
tenant.status = TenantStatus.ARCHIVE
|
||||
data = _make_auth_data(app=app)
|
||||
with patch("controllers.openapi.auth.prepare.TenantService.get_tenant_by_id", return_value=tenant):
|
||||
with pytest.raises(Forbidden):
|
||||
load_tenant(data)
|
||||
|
||||
|
||||
def test_load_tenant_raises_forbidden_when_missing():
|
||||
app = MagicMock()
|
||||
app.tenant_id = uuid.uuid4()
|
||||
data = _make_auth_data(app=app)
|
||||
with patch("controllers.openapi.auth.prepare.TenantService.get_tenant_by_id", return_value=None):
|
||||
with pytest.raises(Forbidden):
|
||||
load_tenant(data)
|
||||
|
||||
|
||||
def test_load_tenant_raises_500_when_app_not_loaded():
|
||||
from werkzeug.exceptions import InternalServerError
|
||||
|
||||
data = _make_auth_data()
|
||||
with pytest.raises(InternalServerError):
|
||||
load_tenant(data)
|
||||
|
||||
|
||||
def test_load_account_writes_caller():
|
||||
account = MagicMock()
|
||||
account_id = uuid.uuid4()
|
||||
data = _make_auth_data(account_id=account_id)
|
||||
with patch("controllers.openapi.auth.prepare.AccountService.get_account_by_id", return_value=account):
|
||||
load_account(data)
|
||||
assert data.caller is account
|
||||
assert data.caller_kind == "account"
|
||||
|
||||
|
||||
def test_load_account_skips_when_already_set():
|
||||
existing_caller = MagicMock()
|
||||
data = _make_auth_data(account_id=uuid.uuid4(), caller=existing_caller)
|
||||
load_account(data)
|
||||
assert data.caller is existing_caller
|
||||
|
||||
|
||||
def test_load_account_sets_current_tenant_when_tenant_present(sqlite_session: Session):
|
||||
account = MagicMock()
|
||||
tenant = MagicMock()
|
||||
session = sqlite_session
|
||||
data = _make_auth_data(account_id=uuid.uuid4(), tenant=tenant)
|
||||
with (
|
||||
patch("controllers.openapi.auth.prepare.AccountService.get_account_by_id", return_value=account),
|
||||
patch("controllers.openapi.auth.prepare.session_factory.create_session", return_value=nullcontext(session)),
|
||||
):
|
||||
load_account(data)
|
||||
account.set_current_tenant_with_session.assert_called_once_with(tenant, session=session)
|
||||
|
||||
|
||||
def test_load_account_raises_unauthorized_when_not_found():
|
||||
data = _make_auth_data(account_id=uuid.uuid4())
|
||||
with patch("controllers.openapi.auth.prepare.AccountService.get_account_by_id", return_value=None):
|
||||
with pytest.raises(Unauthorized):
|
||||
load_account(data)
|
||||
|
||||
|
||||
def test_resolve_external_user_writes_caller():
|
||||
tenant = MagicMock()
|
||||
app = MagicMock()
|
||||
end_user = MagicMock()
|
||||
ext = ExternalIdentity(email="user@sso.com")
|
||||
data = _make_auth_data(tenant=tenant, app=app, external_identity=ext)
|
||||
with patch("controllers.openapi.auth.prepare.EndUserService.get_or_create_end_user_by_type", return_value=end_user):
|
||||
resolve_external_user(data)
|
||||
assert data.caller is end_user
|
||||
assert data.caller_kind == "end_user"
|
||||
|
||||
|
||||
def test_resolve_external_user_raises_unauthorized_when_context_missing():
|
||||
data = _make_auth_data(tenant=None, app=MagicMock(), external_identity=ExternalIdentity(email="u@s.com"))
|
||||
with pytest.raises(Unauthorized):
|
||||
resolve_external_user(data)
|
||||
|
||||
|
||||
def test_load_app_access_mode_writes_mode():
|
||||
from services.enterprise.enterprise_service import WebAppAccessMode
|
||||
|
||||
app = MagicMock()
|
||||
app.id = "app-1"
|
||||
settings = MagicMock()
|
||||
settings.access_mode = "public"
|
||||
data = _make_auth_data(app=app)
|
||||
with patch(
|
||||
"controllers.openapi.auth.prepare.EnterpriseService.WebAppAuth.get_app_access_mode_by_id",
|
||||
return_value=settings,
|
||||
):
|
||||
load_app_access_mode(data)
|
||||
assert data.app_access_mode == WebAppAccessMode.PUBLIC
|
||||
|
||||
|
||||
def test_load_app_access_mode_writes_none_when_value_error():
|
||||
app = MagicMock()
|
||||
app.id = "app-1"
|
||||
data = _make_auth_data(app=app)
|
||||
with patch(
|
||||
"controllers.openapi.auth.prepare.EnterpriseService.WebAppAuth.get_app_access_mode_by_id",
|
||||
side_effect=ValueError("No data found."),
|
||||
):
|
||||
load_app_access_mode(data)
|
||||
assert data.app_access_mode is None
|
||||
|
||||
|
||||
def test_load_app_access_mode_no_op_when_app_missing():
|
||||
data = _make_auth_data()
|
||||
load_app_access_mode(data)
|
||||
assert data.app_access_mode is None
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def flask_app():
|
||||
return Flask(__name__)
|
||||
|
||||
|
||||
def test_load_tenant_from_request_from_path_params(flask_app):
|
||||
tenant = MagicMock()
|
||||
tenant.status = "normal"
|
||||
wid = str(uuid.uuid4())
|
||||
data = _make_auth_data(path_params={"workspace_id": wid})
|
||||
with flask_app.test_request_context("/test"):
|
||||
with patch("controllers.openapi.auth.prepare.TenantService.get_tenant_by_id", return_value=tenant):
|
||||
load_tenant_from_request(data)
|
||||
assert data.tenant is tenant
|
||||
|
||||
|
||||
def test_load_tenant_from_request_from_query_param(flask_app):
|
||||
tenant = MagicMock()
|
||||
tenant.status = "normal"
|
||||
wid = str(uuid.uuid4())
|
||||
data = _make_auth_data(path_params={})
|
||||
with flask_app.test_request_context(f"/test?workspace_id={wid}"):
|
||||
with patch("controllers.openapi.auth.prepare.TenantService.get_tenant_by_id", return_value=tenant):
|
||||
load_tenant_from_request(data)
|
||||
assert data.tenant is tenant
|
||||
|
||||
|
||||
def test_load_tenant_from_request_skips_when_already_set(flask_app):
|
||||
existing_tenant = MagicMock()
|
||||
data = _make_auth_data(tenant=existing_tenant, path_params={})
|
||||
with flask_app.test_request_context("/test"):
|
||||
load_tenant_from_request(data)
|
||||
assert data.tenant is existing_tenant
|
||||
|
||||
|
||||
def test_load_tenant_from_request_raises_not_found_when_no_id(flask_app):
|
||||
data = _make_auth_data(path_params={})
|
||||
with flask_app.test_request_context("/test"):
|
||||
with pytest.raises(NotFound):
|
||||
load_tenant_from_request(data)
|
||||
|
||||
|
||||
def test_load_tenant_from_request_raises_not_found_when_missing(flask_app):
|
||||
wid = str(uuid.uuid4())
|
||||
data = _make_auth_data(path_params={"workspace_id": wid})
|
||||
with flask_app.test_request_context("/test"):
|
||||
with patch("controllers.openapi.auth.prepare.TenantService.get_tenant_by_id", return_value=None):
|
||||
with pytest.raises(NotFound):
|
||||
load_tenant_from_request(data)
|
||||
|
||||
|
||||
def test_load_tenant_from_request_raises_not_found_when_archived(flask_app):
|
||||
from models.account import TenantStatus
|
||||
|
||||
tenant = MagicMock()
|
||||
tenant.status = TenantStatus.ARCHIVE
|
||||
wid = str(uuid.uuid4())
|
||||
data = _make_auth_data(path_params={"workspace_id": wid})
|
||||
with flask_app.test_request_context("/test"):
|
||||
with patch("controllers.openapi.auth.prepare.TenantService.get_tenant_by_id", return_value=tenant):
|
||||
with pytest.raises(NotFound):
|
||||
load_tenant_from_request(data)
|
||||
|
||||
|
||||
def test_load_tenant_from_request_raises_not_found_when_invalid_uuid(flask_app):
|
||||
data = _make_auth_data(path_params={"workspace_id": "not-a-uuid"})
|
||||
with flask_app.test_request_context("/test"):
|
||||
with pytest.raises(NotFound):
|
||||
load_tenant_from_request(data)
|
||||
|
||||
|
||||
# --- load_workspace_role ---
|
||||
|
||||
|
||||
def test_load_workspace_role_stashes_role():
|
||||
tenant = MagicMock()
|
||||
tenant.id = uuid.uuid4()
|
||||
caller = MagicMock()
|
||||
caller.status = "active"
|
||||
data = _make_auth_data(account_id=uuid.uuid4(), tenant=tenant, caller=caller)
|
||||
with patch(
|
||||
"controllers.openapi.auth.prepare.TenantService.get_account_role_in_tenant",
|
||||
return_value=TenantAccountRole.ADMIN,
|
||||
):
|
||||
load_workspace_role(data)
|
||||
assert data.tenant_role == TenantAccountRole.ADMIN
|
||||
|
||||
|
||||
def test_load_workspace_role_none_when_not_member():
|
||||
tenant = MagicMock()
|
||||
tenant.id = uuid.uuid4()
|
||||
caller = MagicMock()
|
||||
caller.status = "active"
|
||||
data = _make_auth_data(account_id=uuid.uuid4(), tenant=tenant, caller=caller)
|
||||
with patch(
|
||||
"controllers.openapi.auth.prepare.TenantService.get_account_role_in_tenant",
|
||||
return_value=None,
|
||||
):
|
||||
load_workspace_role(data)
|
||||
assert data.tenant_role is None
|
||||
|
||||
|
||||
def test_load_workspace_role_none_when_account_inactive():
|
||||
tenant = MagicMock()
|
||||
tenant.id = uuid.uuid4()
|
||||
caller = MagicMock()
|
||||
caller.status = "banned"
|
||||
data = _make_auth_data(account_id=uuid.uuid4(), tenant=tenant, caller=caller)
|
||||
load_workspace_role(data)
|
||||
assert data.tenant_role is None
|
||||
|
||||
|
||||
def test_load_workspace_role_skips_when_already_set():
|
||||
tenant = MagicMock()
|
||||
tenant.id = uuid.uuid4()
|
||||
caller = MagicMock()
|
||||
caller.status = "active"
|
||||
data = _make_auth_data(
|
||||
account_id=uuid.uuid4(),
|
||||
tenant=tenant,
|
||||
caller=caller,
|
||||
tenant_role=TenantAccountRole.OWNER,
|
||||
def _app(
|
||||
*,
|
||||
app_id: str = APP_ID,
|
||||
tenant_id: str = TENANT_ID,
|
||||
enable_api: bool = True,
|
||||
) -> App:
|
||||
return App(
|
||||
id=app_id,
|
||||
tenant_id=tenant_id,
|
||||
name="OpenAPI app",
|
||||
description="",
|
||||
mode=AppMode.CHAT,
|
||||
icon_type=IconType.EMOJI,
|
||||
icon="robot",
|
||||
icon_background="#FFFFFF",
|
||||
status=AppStatus.NORMAL,
|
||||
enable_site=True,
|
||||
enable_api=enable_api,
|
||||
max_active_requests=None,
|
||||
)
|
||||
load_workspace_role(data)
|
||||
assert data.tenant_role == TenantAccountRole.OWNER
|
||||
|
||||
|
||||
def test_load_workspace_role_skips_when_tenant_missing():
|
||||
data = _make_auth_data(account_id=uuid.uuid4())
|
||||
load_workspace_role(data)
|
||||
assert data.tenant_role is None
|
||||
def _tenant(*, tenant_id: str = TENANT_ID, status: TenantStatus = TenantStatus.NORMAL) -> Tenant:
|
||||
tenant = Tenant(name="OpenAPI tenant", status=status)
|
||||
tenant.id = tenant_id
|
||||
return tenant
|
||||
|
||||
|
||||
def test_load_workspace_role_skips_when_account_id_missing():
|
||||
tenant = MagicMock()
|
||||
data = _make_auth_data(tenant=tenant, account_id=None)
|
||||
load_workspace_role(data)
|
||||
assert data.tenant_role is None
|
||||
def _account(*, status: AccountStatus = AccountStatus.ACTIVE) -> Account:
|
||||
account = Account(name="OpenAPI account", email="account@example.com", status=status)
|
||||
account.id = ACCOUNT_ID
|
||||
return account
|
||||
|
||||
|
||||
def _persist(session: Session, *models: object) -> None:
|
||||
session.add_all(models)
|
||||
session.commit()
|
||||
|
||||
|
||||
class TestLoadApp:
|
||||
def test_writes_persisted_app_to_data(self, sqlite_session: Session) -> None:
|
||||
_persist(sqlite_session, _app())
|
||||
data = _make_auth_data(path_params={"app_id": APP_ID})
|
||||
|
||||
load_app(data)
|
||||
|
||||
assert data.app is not None
|
||||
assert data.app.id == APP_ID
|
||||
|
||||
def test_rejects_non_uuid_and_missing_app(self) -> None:
|
||||
with pytest.raises(NotFound, match="app not found"):
|
||||
load_app(_make_auth_data(path_params={"app_id": "not-a-uuid"}))
|
||||
with pytest.raises(NotFound, match="app not found"):
|
||||
load_app(_make_auth_data(path_params={"app_id": APP_ID}))
|
||||
|
||||
def test_rejects_non_normal_app(self) -> None:
|
||||
app = _app()
|
||||
app.status = "archived" # type: ignore[assignment]
|
||||
|
||||
with (
|
||||
patch("controllers.openapi.auth.prepare.AppService.get_app_by_id", return_value=app),
|
||||
pytest.raises(NotFound, match="app not found"),
|
||||
):
|
||||
load_app(_make_auth_data(path_params={"app_id": APP_ID}))
|
||||
|
||||
def test_stashes_app_even_when_api_disabled(self, sqlite_session: Session) -> None:
|
||||
_persist(sqlite_session, _app(enable_api=False))
|
||||
data = _make_auth_data(path_params={"app_id": APP_ID})
|
||||
|
||||
load_app(data)
|
||||
|
||||
assert data.app is not None
|
||||
assert data.app.enable_api is False
|
||||
|
||||
def test_skips_when_already_set(self) -> None:
|
||||
existing_app = _app()
|
||||
data = _make_auth_data(app=existing_app, path_params={"app_id": "invalid"})
|
||||
|
||||
load_app(data)
|
||||
|
||||
assert data.app is existing_app
|
||||
|
||||
|
||||
class TestLoadTenant:
|
||||
def test_writes_persisted_tenant(self, sqlite_session: Session) -> None:
|
||||
app = _app()
|
||||
_persist(sqlite_session, app, _tenant())
|
||||
data = _make_auth_data(app=app)
|
||||
|
||||
load_tenant(data)
|
||||
|
||||
assert data.tenant is not None
|
||||
assert data.tenant.id == TENANT_ID
|
||||
|
||||
def test_skips_when_already_set(self) -> None:
|
||||
tenant = _tenant()
|
||||
data = _make_auth_data(app=_app(), tenant=tenant)
|
||||
|
||||
load_tenant(data)
|
||||
|
||||
assert data.tenant is tenant
|
||||
|
||||
@pytest.mark.parametrize("persist_archived", [True, False])
|
||||
def test_rejects_archived_or_missing_tenant(self, sqlite_session: Session, persist_archived: bool) -> None:
|
||||
app = _app()
|
||||
models: list[object] = [app]
|
||||
if persist_archived:
|
||||
models.append(_tenant(status=TenantStatus.ARCHIVE))
|
||||
_persist(sqlite_session, *models)
|
||||
|
||||
with pytest.raises(Forbidden, match="workspace unavailable"):
|
||||
load_tenant(_make_auth_data(app=app))
|
||||
|
||||
def test_rejects_missing_app_context(self) -> None:
|
||||
with pytest.raises(InternalServerError, match="app not loaded"):
|
||||
load_tenant(_make_auth_data())
|
||||
|
||||
|
||||
class TestLoadAccount:
|
||||
def test_writes_persisted_caller(self, sqlite_session: Session) -> None:
|
||||
_persist(sqlite_session, _account())
|
||||
data = _make_auth_data(account_id=uuid.UUID(ACCOUNT_ID))
|
||||
|
||||
load_account(data)
|
||||
|
||||
assert data.caller is not None
|
||||
assert data.caller.id == ACCOUNT_ID
|
||||
assert data.caller_kind == "account"
|
||||
|
||||
def test_sets_current_tenant_from_real_membership(self, sqlite_session: Session) -> None:
|
||||
account = _account()
|
||||
tenant = _tenant()
|
||||
membership = TenantAccountJoin(
|
||||
tenant_id=tenant.id,
|
||||
account_id=account.id,
|
||||
current=True,
|
||||
role=TenantAccountRole.ADMIN,
|
||||
)
|
||||
_persist(sqlite_session, account, tenant, membership)
|
||||
data = _make_auth_data(account_id=uuid.UUID(ACCOUNT_ID), tenant=tenant)
|
||||
|
||||
load_account(data)
|
||||
|
||||
assert isinstance(data.caller, Account)
|
||||
assert data.caller.current_tenant_id == TENANT_ID
|
||||
assert data.caller.role == TenantAccountRole.ADMIN
|
||||
|
||||
def test_skips_when_caller_already_set(self) -> None:
|
||||
account = _account()
|
||||
data = _make_auth_data(account_id=uuid.UUID(ACCOUNT_ID), caller=account)
|
||||
|
||||
load_account(data)
|
||||
|
||||
assert data.caller is account
|
||||
|
||||
def test_rejects_missing_account(self) -> None:
|
||||
with pytest.raises(Unauthorized, match="account not found"):
|
||||
load_account(_make_auth_data(account_id=uuid.UUID(ACCOUNT_ID)))
|
||||
|
||||
|
||||
class TestResolveExternalUser:
|
||||
def test_persists_and_writes_end_user(
|
||||
self,
|
||||
sqlite_engine: Engine,
|
||||
sqlite_session: Session,
|
||||
) -> None:
|
||||
app = _app()
|
||||
tenant = _tenant()
|
||||
_persist(sqlite_session, app, tenant)
|
||||
data = _make_auth_data(
|
||||
tenant=tenant,
|
||||
app=app,
|
||||
external_identity=ExternalIdentity(email="user@sso.com"),
|
||||
)
|
||||
|
||||
with patch.object(type(end_user_service.db), "engine", new_callable=PropertyMock) as engine:
|
||||
engine.return_value = sqlite_engine
|
||||
resolve_external_user(data)
|
||||
|
||||
assert isinstance(data.caller, EndUser)
|
||||
assert data.caller_kind == "end_user"
|
||||
with Session(sqlite_engine) as observer:
|
||||
persisted = observer.scalar(select(EndUser).where(EndUser.session_id == "user@sso.com"))
|
||||
assert persisted is not None
|
||||
assert persisted.tenant_id == TENANT_ID
|
||||
assert persisted.app_id == APP_ID
|
||||
|
||||
def test_rejects_missing_context(self) -> None:
|
||||
data = _make_auth_data(app=_app(), external_identity=ExternalIdentity(email="u@s.com"))
|
||||
|
||||
with pytest.raises(Unauthorized, match="missing context"):
|
||||
resolve_external_user(data)
|
||||
|
||||
|
||||
class TestLoadAppAccessMode:
|
||||
def test_writes_mode(self) -> None:
|
||||
data = _make_auth_data(app=_app())
|
||||
settings = SimpleNamespace(access_mode="public")
|
||||
|
||||
with patch(
|
||||
"controllers.openapi.auth.prepare.EnterpriseService.WebAppAuth.get_app_access_mode_by_id",
|
||||
return_value=settings,
|
||||
):
|
||||
load_app_access_mode(data)
|
||||
|
||||
assert data.app_access_mode == WebAppAccessMode.PUBLIC
|
||||
|
||||
def test_writes_none_when_provider_raises(self) -> None:
|
||||
data = _make_auth_data(app=_app())
|
||||
with patch(
|
||||
"controllers.openapi.auth.prepare.EnterpriseService.WebAppAuth.get_app_access_mode_by_id",
|
||||
side_effect=ValueError("No data found."),
|
||||
):
|
||||
load_app_access_mode(data)
|
||||
assert data.app_access_mode is None
|
||||
|
||||
def test_noop_without_app(self) -> None:
|
||||
data = _make_auth_data()
|
||||
load_app_access_mode(data)
|
||||
assert data.app_access_mode is None
|
||||
|
||||
|
||||
class TestLoadTenantFromRequest:
|
||||
def test_loads_from_path_or_query(
|
||||
self,
|
||||
app: Flask,
|
||||
sqlite_session: Session,
|
||||
) -> None:
|
||||
_persist(sqlite_session, _tenant())
|
||||
for path_params, path in (({"workspace_id": TENANT_ID}, "/test"), ({}, f"/test?workspace_id={TENANT_ID}")):
|
||||
data = _make_auth_data(path_params=path_params)
|
||||
with app.test_request_context(path):
|
||||
load_tenant_from_request(data)
|
||||
assert data.tenant is not None
|
||||
assert data.tenant.id == TENANT_ID
|
||||
|
||||
def test_skips_when_already_set(self, app: Flask) -> None:
|
||||
tenant = _tenant()
|
||||
data = _make_auth_data(tenant=tenant)
|
||||
with app.test_request_context("/test"):
|
||||
load_tenant_from_request(data)
|
||||
assert data.tenant is tenant
|
||||
|
||||
def test_rejects_missing_or_invalid_id(self, app: Flask) -> None:
|
||||
for path_params in ({}, {"workspace_id": "not-a-uuid"}):
|
||||
with app.test_request_context("/test"), pytest.raises(NotFound, match="workspace not found"):
|
||||
load_tenant_from_request(_make_auth_data(path_params=path_params))
|
||||
|
||||
@pytest.mark.parametrize("tenant_status", [None, TenantStatus.ARCHIVE])
|
||||
def test_rejects_missing_or_archived_tenant(
|
||||
self,
|
||||
app: Flask,
|
||||
sqlite_session: Session,
|
||||
tenant_status: TenantStatus | None,
|
||||
) -> None:
|
||||
if tenant_status is not None:
|
||||
_persist(sqlite_session, _tenant(status=tenant_status))
|
||||
data = _make_auth_data(path_params={"workspace_id": TENANT_ID})
|
||||
|
||||
with app.test_request_context("/test"), pytest.raises(NotFound, match="workspace not found"):
|
||||
load_tenant_from_request(data)
|
||||
|
||||
|
||||
class TestLoadWorkspaceRole:
|
||||
def test_loads_real_membership_role(self, sqlite_session: Session) -> None:
|
||||
account = _account()
|
||||
tenant = _tenant()
|
||||
membership = TenantAccountJoin(
|
||||
tenant_id=tenant.id,
|
||||
account_id=account.id,
|
||||
current=True,
|
||||
role=TenantAccountRole.ADMIN,
|
||||
)
|
||||
_persist(sqlite_session, account, tenant, membership)
|
||||
data = _make_auth_data(
|
||||
account_id=uuid.UUID(ACCOUNT_ID),
|
||||
tenant=tenant,
|
||||
caller=account,
|
||||
)
|
||||
|
||||
load_workspace_role(data)
|
||||
|
||||
assert data.tenant_role == TenantAccountRole.ADMIN
|
||||
|
||||
def test_none_when_not_member(self, sqlite_session: Session) -> None:
|
||||
account = _account()
|
||||
tenant = _tenant()
|
||||
_persist(sqlite_session, account, tenant)
|
||||
data = _make_auth_data(account_id=uuid.UUID(ACCOUNT_ID), tenant=tenant, caller=account)
|
||||
|
||||
load_workspace_role(data)
|
||||
|
||||
assert data.tenant_role is None
|
||||
|
||||
def test_none_when_account_inactive(self) -> None:
|
||||
data = _make_auth_data(
|
||||
account_id=uuid.UUID(ACCOUNT_ID),
|
||||
tenant=_tenant(),
|
||||
caller=_account(status=AccountStatus.BANNED),
|
||||
)
|
||||
load_workspace_role(data)
|
||||
assert data.tenant_role is None
|
||||
|
||||
def test_skips_when_already_set(self) -> None:
|
||||
data = _make_auth_data(
|
||||
account_id=uuid.UUID(ACCOUNT_ID),
|
||||
tenant=_tenant(),
|
||||
caller=_account(),
|
||||
tenant_role=TenantAccountRole.OWNER,
|
||||
)
|
||||
load_workspace_role(data)
|
||||
assert data.tenant_role == TenantAccountRole.OWNER
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"data",
|
||||
[
|
||||
_make_auth_data(account_id=uuid.UUID(ACCOUNT_ID)),
|
||||
_make_auth_data(tenant=_tenant(), account_id=None),
|
||||
],
|
||||
)
|
||||
def test_skips_without_tenant_or_account(self, data: AuthData) -> None:
|
||||
load_workspace_role(data)
|
||||
assert data.tenant_role is None
|
||||
|
||||
@ -4,12 +4,12 @@ Unit tests for Service API wraps (authentication decorators)
|
||||
|
||||
import uuid
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
from flask import Flask
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.orm import Session
|
||||
from sqlalchemy.orm import Session, scoped_session
|
||||
from werkzeug.exceptions import Forbidden, NotFound, ServiceUnavailable, Unauthorized
|
||||
|
||||
from controllers.service_api.wraps import (
|
||||
@ -36,11 +36,9 @@ def _configure_current_app_mock(mock_current_app):
|
||||
mock_current_app._get_current_object = Mock(return_value=Mock())
|
||||
|
||||
|
||||
def _session_proxy(session: Session) -> MagicMock:
|
||||
"""Emulate Flask-SQLAlchemy's callable scoped-session proxy around a test session."""
|
||||
proxy = MagicMock(wraps=session)
|
||||
proxy.return_value = session
|
||||
return proxy
|
||||
def _session_proxy(session: Session) -> scoped_session[Session]:
|
||||
"""Expose the real SQLite session through Flask-SQLAlchemy's callable shape."""
|
||||
return scoped_session(lambda: session)
|
||||
|
||||
|
||||
def _api_token(*, tenant_id: str, app_id: str | None = None, token_type: ApiTokenType) -> ApiToken:
|
||||
|
||||
Loading…
Reference in New Issue
Block a user