Merge remote-tracking branch 'upstream/main' into feat/human-input-merge-again

This commit is contained in:
QuantumGhost
2026-01-28 16:21:37 +08:00
4167 changed files with 345823 additions and 171263 deletions

View File

@ -0,0 +1,69 @@
import builtins
from types import SimpleNamespace
from unittest.mock import patch
from flask.views import MethodView as FlaskMethodView
_NEEDS_METHOD_VIEW_CLEANUP = False
if not hasattr(builtins, "MethodView"):
builtins.MethodView = FlaskMethodView
_NEEDS_METHOD_VIEW_CLEANUP = True
from controllers.common.fields import Parameters, Site
from core.app.app_config.common.parameters_mapping import get_parameters_from_feature_dict
from models.model import IconType
def test_parameters_model_round_trip():
parameters = get_parameters_from_feature_dict(features_dict={}, user_input_form=[])
model = Parameters.model_validate(parameters)
assert model.model_dump(mode="json") == parameters
def test_site_icon_url_uses_signed_url_for_image_icon():
site = SimpleNamespace(
title="Example",
chat_color_theme=None,
chat_color_theme_inverted=False,
icon_type=IconType.IMAGE,
icon="file-id",
icon_background=None,
description=None,
copyright=None,
privacy_policy=None,
custom_disclaimer=None,
default_language="en-US",
show_workflow_steps=True,
use_icon_as_answer_icon=False,
)
with patch("controllers.common.fields.file_helpers.get_signed_file_url", return_value="signed") as mock_helper:
model = Site.model_validate(site)
assert model.icon_url == "signed"
mock_helper.assert_called_once_with("file-id")
def test_site_icon_url_is_none_for_non_image_icon():
site = SimpleNamespace(
title="Example",
chat_color_theme=None,
chat_color_theme_inverted=False,
icon_type=IconType.EMOJI,
icon="file-id",
icon_background=None,
description=None,
copyright=None,
privacy_policy=None,
custom_disclaimer=None,
default_language="en-US",
show_workflow_steps=True,
use_icon_as_answer_icon=False,
)
with patch("controllers.common.fields.file_helpers.get_signed_file_url") as mock_helper:
model = Site.model_validate(site)
assert model.icon_url is None
mock_helper.assert_not_called()

View File

@ -0,0 +1,285 @@
from __future__ import annotations
import builtins
import sys
from datetime import datetime
from importlib import util
from pathlib import Path
from types import ModuleType, SimpleNamespace
from typing import Any
import pytest
from flask.views import MethodView
# kombu references MethodView as a global when importing celery/kombu pools.
if not hasattr(builtins, "MethodView"):
builtins.MethodView = MethodView # type: ignore[attr-defined]
def _load_app_module():
module_name = "controllers.console.app.app"
if module_name in sys.modules:
return sys.modules[module_name]
root = Path(__file__).resolve().parents[5]
module_path = root / "controllers" / "console" / "app" / "app.py"
class _StubNamespace:
def __init__(self):
self.models: dict[str, Any] = {}
self.payload = None
def schema_model(self, name, schema):
self.models[name] = schema
def _decorator(self, obj):
return obj
def doc(self, *args, **kwargs):
return self._decorator
def expect(self, *args, **kwargs):
return self._decorator
def response(self, *args, **kwargs):
return self._decorator
def route(self, *args, **kwargs):
def decorator(obj):
return obj
return decorator
stub_namespace = _StubNamespace()
original_console = sys.modules.get("controllers.console")
original_app_pkg = sys.modules.get("controllers.console.app")
stubbed_modules: list[tuple[str, ModuleType | None]] = []
console_module = ModuleType("controllers.console")
console_module.__path__ = [str(root / "controllers" / "console")]
console_module.console_ns = stub_namespace
console_module.api = None
console_module.bp = None
sys.modules["controllers.console"] = console_module
app_package = ModuleType("controllers.console.app")
app_package.__path__ = [str(root / "controllers" / "console" / "app")]
sys.modules["controllers.console.app"] = app_package
console_module.app = app_package
def _stub_module(name: str, attrs: dict[str, Any]):
original = sys.modules.get(name)
module = ModuleType(name)
for key, value in attrs.items():
setattr(module, key, value)
sys.modules[name] = module
stubbed_modules.append((name, original))
class _OpsTraceManager:
@staticmethod
def get_app_tracing_config(app_id: str) -> dict[str, Any]:
return {}
@staticmethod
def update_app_tracing_config(app_id: str, **kwargs) -> None:
return None
_stub_module(
"core.ops.ops_trace_manager",
{
"OpsTraceManager": _OpsTraceManager,
"TraceQueueManager": object,
"TraceTask": object,
},
)
spec = util.spec_from_file_location(module_name, module_path)
module = util.module_from_spec(spec)
sys.modules[module_name] = module
try:
assert spec.loader is not None
spec.loader.exec_module(module)
finally:
for name, original in reversed(stubbed_modules):
if original is not None:
sys.modules[name] = original
else:
sys.modules.pop(name, None)
if original_console is not None:
sys.modules["controllers.console"] = original_console
else:
sys.modules.pop("controllers.console", None)
if original_app_pkg is not None:
sys.modules["controllers.console.app"] = original_app_pkg
else:
sys.modules.pop("controllers.console.app", None)
return module
_app_module = _load_app_module()
AppDetailWithSite = _app_module.AppDetailWithSite
AppPagination = _app_module.AppPagination
AppPartial = _app_module.AppPartial
@pytest.fixture(autouse=True)
def patch_signed_url(monkeypatch):
"""Ensure icon URL generation uses a deterministic helper for tests."""
def _fake_signed_url(key: str | None) -> str | None:
if not key:
return None
return f"signed:{key}"
monkeypatch.setattr(_app_module.file_helpers, "get_signed_file_url", _fake_signed_url)
def _ts(hour: int = 12) -> datetime:
return datetime(2024, 1, 1, hour, 0, 0)
def _dummy_model_config():
return SimpleNamespace(
model_dict={"provider": "openai", "name": "gpt-4o"},
pre_prompt="hello",
created_by="config-author",
created_at=_ts(9),
updated_by="config-editor",
updated_at=_ts(10),
)
def _dummy_workflow():
return SimpleNamespace(
id="wf-1",
created_by="workflow-author",
created_at=_ts(8),
updated_by="workflow-editor",
updated_at=_ts(9),
)
def test_app_partial_serialization_uses_aliases():
created_at = _ts()
app_obj = SimpleNamespace(
id="app-1",
name="My App",
desc_or_prompt="Prompt snippet",
mode_compatible_with_agent="chat",
icon_type="image",
icon="icon-key",
icon_background="#fff",
app_model_config=_dummy_model_config(),
workflow=_dummy_workflow(),
created_by="creator",
created_at=created_at,
updated_by="editor",
updated_at=created_at,
tags=[SimpleNamespace(id="tag-1", name="Utilities", type="app")],
access_mode="private",
create_user_name="Creator",
author_name="Author",
has_draft_trigger=True,
)
serialized = AppPartial.model_validate(app_obj, from_attributes=True).model_dump(mode="json")
assert serialized["description"] == "Prompt snippet"
assert serialized["mode"] == "chat"
assert serialized["icon_url"] == "signed:icon-key"
assert serialized["created_at"] == int(created_at.timestamp())
assert serialized["updated_at"] == int(created_at.timestamp())
assert serialized["model_config"]["model"] == {"provider": "openai", "name": "gpt-4o"}
assert serialized["workflow"]["id"] == "wf-1"
assert serialized["tags"][0]["name"] == "Utilities"
def test_app_detail_with_site_includes_nested_serialization():
timestamp = _ts(14)
site = SimpleNamespace(
code="site-code",
title="Public Site",
icon_type="image",
icon="site-icon",
created_at=timestamp,
updated_at=timestamp,
)
app_obj = SimpleNamespace(
id="app-2",
name="Detailed App",
description="Desc",
mode_compatible_with_agent="advanced-chat",
icon_type="image",
icon="detail-icon",
icon_background="#123456",
enable_site=True,
enable_api=True,
app_model_config={
"opening_statement": "hi",
"model": {"provider": "openai", "name": "gpt-4o"},
"retriever_resource": {"enabled": True},
},
workflow=_dummy_workflow(),
tracing={"enabled": True},
use_icon_as_answer_icon=True,
created_by="creator",
created_at=timestamp,
updated_by="editor",
updated_at=timestamp,
access_mode="public",
tags=[SimpleNamespace(id="tag-2", name="Prod", type="app")],
api_base_url="https://api.example.com/v1",
max_active_requests=5,
deleted_tools=[{"type": "api", "tool_name": "search", "provider_id": "prov"}],
site=site,
)
serialized = AppDetailWithSite.model_validate(app_obj, from_attributes=True).model_dump(mode="json")
assert serialized["icon_url"] == "signed:detail-icon"
assert serialized["model_config"]["retriever_resource"] == {"enabled": True}
assert serialized["deleted_tools"][0]["tool_name"] == "search"
assert serialized["site"]["icon_url"] == "signed:site-icon"
assert serialized["site"]["created_at"] == int(timestamp.timestamp())
def test_app_pagination_aliases_per_page_and_has_next():
item_one = SimpleNamespace(
id="app-10",
name="Paginated One",
desc_or_prompt="Summary",
mode_compatible_with_agent="chat",
icon_type="image",
icon="first-icon",
created_at=_ts(15),
updated_at=_ts(15),
)
item_two = SimpleNamespace(
id="app-11",
name="Paginated Two",
desc_or_prompt="Summary",
mode_compatible_with_agent="agent-chat",
icon_type="emoji",
icon="🙂",
created_at=_ts(16),
updated_at=_ts(16),
)
pagination = SimpleNamespace(
page=2,
per_page=10,
total=50,
has_next=True,
items=[item_one, item_two],
)
serialized = AppPagination.model_validate(pagination, from_attributes=True).model_dump(mode="json")
assert serialized["page"] == 2
assert serialized["limit"] == 10
assert serialized["has_more"] is True
assert len(serialized["data"]) == 2
assert serialized["data"][0]["icon_url"] == "signed:first-icon"
assert serialized["data"][1]["icon_url"] is None

View File

@ -40,7 +40,7 @@ class TestActivateCheckApi:
"tenant": tenant,
}
@patch("controllers.console.auth.activate.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.activate.RegisterService.get_invitation_with_case_fallback")
def test_check_valid_invitation_token(self, mock_get_invitation, app, mock_invitation):
"""
Test checking valid invitation token.
@ -66,7 +66,7 @@ class TestActivateCheckApi:
assert response["data"]["workspace_id"] == "workspace-123"
assert response["data"]["email"] == "invitee@example.com"
@patch("controllers.console.auth.activate.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.activate.RegisterService.get_invitation_with_case_fallback")
def test_check_invalid_invitation_token(self, mock_get_invitation, app):
"""
Test checking invalid invitation token.
@ -88,7 +88,7 @@ class TestActivateCheckApi:
# Assert
assert response["is_valid"] is False
@patch("controllers.console.auth.activate.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.activate.RegisterService.get_invitation_with_case_fallback")
def test_check_token_without_workspace_id(self, mock_get_invitation, app, mock_invitation):
"""
Test checking token without workspace ID.
@ -109,7 +109,7 @@ class TestActivateCheckApi:
assert response["is_valid"] is True
mock_get_invitation.assert_called_once_with(None, "invitee@example.com", "valid_token")
@patch("controllers.console.auth.activate.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.activate.RegisterService.get_invitation_with_case_fallback")
def test_check_token_without_email(self, mock_get_invitation, app, mock_invitation):
"""
Test checking token without email parameter.
@ -130,6 +130,20 @@ class TestActivateCheckApi:
assert response["is_valid"] is True
mock_get_invitation.assert_called_once_with("workspace-123", None, "valid_token")
@patch("controllers.console.auth.activate.RegisterService.get_invitation_with_case_fallback")
def test_check_token_normalizes_email_to_lowercase(self, mock_get_invitation, app, mock_invitation):
"""Ensure token validation uses lowercase emails."""
mock_get_invitation.return_value = mock_invitation
with app.test_request_context(
"/activate/check?workspace_id=workspace-123&email=Invitee@Example.com&token=valid_token"
):
api = ActivateCheckApi()
response = api.get()
assert response["is_valid"] is True
mock_get_invitation.assert_called_once_with("workspace-123", "Invitee@Example.com", "valid_token")
class TestActivateApi:
"""Test cases for account activation endpoint."""
@ -212,7 +226,7 @@ class TestActivateApi:
mock_revoke_token.assert_called_once_with("workspace-123", "invitee@example.com", "valid_token")
mock_db.session.commit.assert_called_once()
@patch("controllers.console.auth.activate.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.activate.RegisterService.get_invitation_with_case_fallback")
def test_activation_with_invalid_token(self, mock_get_invitation, app):
"""
Test account activation with invalid token.
@ -241,7 +255,7 @@ class TestActivateApi:
with pytest.raises(AlreadyActivateError):
api.post()
@patch("controllers.console.auth.activate.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.activate.RegisterService.get_invitation_with_case_fallback")
@patch("controllers.console.auth.activate.RegisterService.revoke_token")
@patch("controllers.console.auth.activate.db")
def test_activation_sets_interface_theme(
@ -290,7 +304,7 @@ class TestActivateApi:
("es-ES", "Europe/Madrid"),
],
)
@patch("controllers.console.auth.activate.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.activate.RegisterService.get_invitation_with_case_fallback")
@patch("controllers.console.auth.activate.RegisterService.revoke_token")
@patch("controllers.console.auth.activate.db")
def test_activation_with_different_locales(
@ -336,7 +350,7 @@ class TestActivateApi:
assert mock_account.interface_language == language
assert mock_account.timezone == timezone
@patch("controllers.console.auth.activate.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.activate.RegisterService.get_invitation_with_case_fallback")
@patch("controllers.console.auth.activate.RegisterService.revoke_token")
@patch("controllers.console.auth.activate.db")
def test_activation_returns_success_response(
@ -376,7 +390,7 @@ class TestActivateApi:
# Assert
assert response == {"result": "success"}
@patch("controllers.console.auth.activate.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.activate.RegisterService.get_invitation_with_case_fallback")
@patch("controllers.console.auth.activate.RegisterService.revoke_token")
@patch("controllers.console.auth.activate.db")
def test_activation_without_workspace_id(
@ -415,3 +429,37 @@ class TestActivateApi:
# Assert
assert response["result"] == "success"
mock_revoke_token.assert_called_once_with(None, "invitee@example.com", "valid_token")
@patch("controllers.console.auth.activate.RegisterService.get_invitation_with_case_fallback")
@patch("controllers.console.auth.activate.RegisterService.revoke_token")
@patch("controllers.console.auth.activate.db")
def test_activation_normalizes_email_before_lookup(
self,
mock_db,
mock_revoke_token,
mock_get_invitation,
app,
mock_invitation,
mock_account,
):
"""Ensure uppercase emails are normalized before lookup and revocation."""
mock_get_invitation.return_value = mock_invitation
with app.test_request_context(
"/activate",
method="POST",
json={
"workspace_id": "workspace-123",
"email": "Invitee@Example.com",
"token": "valid_token",
"name": "John Doe",
"interface_language": "en-US",
"timezone": "UTC",
},
):
api = ActivateApi()
response = api.post()
assert response["result"] == "success"
mock_get_invitation.assert_called_once_with("workspace-123", "Invitee@Example.com", "valid_token")
mock_revoke_token.assert_called_once_with("workspace-123", "invitee@example.com", "valid_token")

View File

@ -34,7 +34,7 @@ class TestAuthenticationSecurity:
@patch("controllers.console.auth.login.AccountService.authenticate")
@patch("controllers.console.auth.login.AccountService.add_login_error_rate_limit")
@patch("controllers.console.auth.login.dify_config.BILLING_ENABLED", False)
@patch("controllers.console.auth.login.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.login.RegisterService.get_invitation_with_case_fallback")
def test_login_invalid_email_with_registration_allowed(
self, mock_get_invitation, mock_add_rate_limit, mock_authenticate, mock_is_rate_limit, mock_features, mock_db
):
@ -67,7 +67,7 @@ class TestAuthenticationSecurity:
@patch("controllers.console.auth.login.AccountService.authenticate")
@patch("controllers.console.auth.login.AccountService.add_login_error_rate_limit")
@patch("controllers.console.auth.login.dify_config.BILLING_ENABLED", False)
@patch("controllers.console.auth.login.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.login.RegisterService.get_invitation_with_case_fallback")
def test_login_wrong_password_returns_error(
self, mock_get_invitation, mock_add_rate_limit, mock_authenticate, mock_is_rate_limit, mock_db
):
@ -100,7 +100,7 @@ class TestAuthenticationSecurity:
@patch("controllers.console.auth.login.AccountService.authenticate")
@patch("controllers.console.auth.login.AccountService.add_login_error_rate_limit")
@patch("controllers.console.auth.login.dify_config.BILLING_ENABLED", False)
@patch("controllers.console.auth.login.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.login.RegisterService.get_invitation_with_case_fallback")
def test_login_invalid_email_with_registration_disabled(
self, mock_get_invitation, mock_add_rate_limit, mock_authenticate, mock_is_rate_limit, mock_features, mock_db
):

View File

@ -0,0 +1,177 @@
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from flask import Flask
from controllers.console.auth.email_register import (
EmailRegisterCheckApi,
EmailRegisterResetApi,
EmailRegisterSendEmailApi,
)
from services.account_service import AccountService
@pytest.fixture
def app():
flask_app = Flask(__name__)
flask_app.config["TESTING"] = True
return flask_app
class TestEmailRegisterSendEmailApi:
@patch("controllers.console.auth.email_register.Session")
@patch("controllers.console.auth.email_register.AccountService.get_account_by_email_with_case_fallback")
@patch("controllers.console.auth.email_register.AccountService.send_email_register_email")
@patch("controllers.console.auth.email_register.BillingService.is_email_in_freeze")
@patch("controllers.console.auth.email_register.AccountService.is_email_send_ip_limit", return_value=False)
@patch("controllers.console.auth.email_register.extract_remote_ip", return_value="127.0.0.1")
def test_send_email_normalizes_and_falls_back(
self,
mock_extract_ip,
mock_is_email_send_ip_limit,
mock_is_freeze,
mock_send_mail,
mock_get_account,
mock_session_cls,
app,
):
mock_send_mail.return_value = "token-123"
mock_is_freeze.return_value = False
mock_account = MagicMock()
mock_session = MagicMock()
mock_session_cls.return_value.__enter__.return_value = mock_session
mock_get_account.return_value = mock_account
feature_flags = SimpleNamespace(enable_email_password_login=True, is_allow_register=True)
with (
patch("controllers.console.auth.email_register.db", SimpleNamespace(engine="engine")),
patch("controllers.console.auth.email_register.dify_config", SimpleNamespace(BILLING_ENABLED=True)),
patch("controllers.console.wraps.dify_config", SimpleNamespace(EDITION="CLOUD")),
patch("controllers.console.wraps.FeatureService.get_system_features", return_value=feature_flags),
):
with app.test_request_context(
"/email-register/send-email",
method="POST",
json={"email": "Invitee@Example.com", "language": "en-US"},
):
response = EmailRegisterSendEmailApi().post()
assert response == {"result": "success", "data": "token-123"}
mock_is_freeze.assert_called_once_with("invitee@example.com")
mock_send_mail.assert_called_once_with(email="invitee@example.com", account=mock_account, language="en-US")
mock_get_account.assert_called_once_with("Invitee@Example.com", session=mock_session)
mock_extract_ip.assert_called_once()
mock_is_email_send_ip_limit.assert_called_once_with("127.0.0.1")
class TestEmailRegisterCheckApi:
@patch("controllers.console.auth.email_register.AccountService.reset_email_register_error_rate_limit")
@patch("controllers.console.auth.email_register.AccountService.generate_email_register_token")
@patch("controllers.console.auth.email_register.AccountService.revoke_email_register_token")
@patch("controllers.console.auth.email_register.AccountService.add_email_register_error_rate_limit")
@patch("controllers.console.auth.email_register.AccountService.get_email_register_data")
@patch("controllers.console.auth.email_register.AccountService.is_email_register_error_rate_limit")
def test_validity_normalizes_email_before_checks(
self,
mock_rate_limit_check,
mock_get_data,
mock_add_rate,
mock_revoke,
mock_generate_token,
mock_reset_rate,
app,
):
mock_rate_limit_check.return_value = False
mock_get_data.return_value = {"email": "User@Example.com", "code": "4321"}
mock_generate_token.return_value = (None, "new-token")
feature_flags = SimpleNamespace(enable_email_password_login=True, is_allow_register=True)
with (
patch("controllers.console.auth.email_register.db", SimpleNamespace(engine="engine")),
patch("controllers.console.wraps.dify_config", SimpleNamespace(EDITION="CLOUD")),
patch("controllers.console.wraps.FeatureService.get_system_features", return_value=feature_flags),
):
with app.test_request_context(
"/email-register/validity",
method="POST",
json={"email": "User@Example.com", "code": "4321", "token": "token-123"},
):
response = EmailRegisterCheckApi().post()
assert response == {"is_valid": True, "email": "user@example.com", "token": "new-token"}
mock_rate_limit_check.assert_called_once_with("user@example.com")
mock_generate_token.assert_called_once_with(
"user@example.com", code="4321", additional_data={"phase": "register"}
)
mock_reset_rate.assert_called_once_with("user@example.com")
mock_add_rate.assert_not_called()
mock_revoke.assert_called_once_with("token-123")
class TestEmailRegisterResetApi:
@patch("controllers.console.auth.email_register.AccountService.reset_login_error_rate_limit")
@patch("controllers.console.auth.email_register.AccountService.login")
@patch("controllers.console.auth.email_register.EmailRegisterResetApi._create_new_account")
@patch("controllers.console.auth.email_register.Session")
@patch("controllers.console.auth.email_register.AccountService.get_account_by_email_with_case_fallback")
@patch("controllers.console.auth.email_register.AccountService.revoke_email_register_token")
@patch("controllers.console.auth.email_register.AccountService.get_email_register_data")
@patch("controllers.console.auth.email_register.extract_remote_ip", return_value="127.0.0.1")
def test_reset_creates_account_with_normalized_email(
self,
mock_extract_ip,
mock_get_data,
mock_revoke_token,
mock_get_account,
mock_session_cls,
mock_create_account,
mock_login,
mock_reset_login_rate,
app,
):
mock_get_data.return_value = {"phase": "register", "email": "Invitee@Example.com"}
mock_create_account.return_value = MagicMock()
token_pair = MagicMock()
token_pair.model_dump.return_value = {"access_token": "a", "refresh_token": "r"}
mock_login.return_value = token_pair
mock_session = MagicMock()
mock_session_cls.return_value.__enter__.return_value = mock_session
mock_get_account.return_value = None
feature_flags = SimpleNamespace(enable_email_password_login=True, is_allow_register=True)
with (
patch("controllers.console.auth.email_register.db", SimpleNamespace(engine="engine")),
patch("controllers.console.wraps.dify_config", SimpleNamespace(EDITION="CLOUD")),
patch("controllers.console.wraps.FeatureService.get_system_features", return_value=feature_flags),
):
with app.test_request_context(
"/email-register",
method="POST",
json={"token": "token-123", "new_password": "ValidPass123!", "password_confirm": "ValidPass123!"},
):
response = EmailRegisterResetApi().post()
assert response == {"result": "success", "data": {"access_token": "a", "refresh_token": "r"}}
mock_create_account.assert_called_once_with("invitee@example.com", "ValidPass123!")
mock_reset_login_rate.assert_called_once_with("invitee@example.com")
mock_revoke_token.assert_called_once_with("token-123")
mock_extract_ip.assert_called_once()
mock_get_account.assert_called_once_with("Invitee@Example.com", session=mock_session)
def test_get_account_by_email_with_case_fallback_uses_lowercase_lookup():
mock_session = MagicMock()
first_query = MagicMock()
first_query.scalar_one_or_none.return_value = None
expected_account = MagicMock()
second_query = MagicMock()
second_query.scalar_one_or_none.return_value = expected_account
mock_session.execute.side_effect = [first_query, second_query]
account = AccountService.get_account_by_email_with_case_fallback("Case@Test.com", session=mock_session)
assert account is expected_account
assert mock_session.execute.call_count == 2

View File

@ -0,0 +1,176 @@
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from flask import Flask
from controllers.console.auth.forgot_password import (
ForgotPasswordCheckApi,
ForgotPasswordResetApi,
ForgotPasswordSendEmailApi,
)
from services.account_service import AccountService
@pytest.fixture
def app():
flask_app = Flask(__name__)
flask_app.config["TESTING"] = True
return flask_app
class TestForgotPasswordSendEmailApi:
@patch("controllers.console.auth.forgot_password.Session")
@patch("controllers.console.auth.forgot_password.AccountService.get_account_by_email_with_case_fallback")
@patch("controllers.console.auth.forgot_password.AccountService.send_reset_password_email")
@patch("controllers.console.auth.forgot_password.AccountService.is_email_send_ip_limit", return_value=False)
@patch("controllers.console.auth.forgot_password.extract_remote_ip", return_value="127.0.0.1")
def test_send_normalizes_email(
self,
mock_extract_ip,
mock_is_ip_limit,
mock_send_email,
mock_get_account,
mock_session_cls,
app,
):
mock_account = MagicMock()
mock_get_account.return_value = mock_account
mock_send_email.return_value = "token-123"
mock_session = MagicMock()
mock_session_cls.return_value.__enter__.return_value = mock_session
wraps_features = SimpleNamespace(enable_email_password_login=True, is_allow_register=True)
controller_features = SimpleNamespace(is_allow_register=True)
with (
patch("controllers.console.auth.forgot_password.db", SimpleNamespace(engine="engine")),
patch(
"controllers.console.auth.forgot_password.FeatureService.get_system_features",
return_value=controller_features,
),
patch("controllers.console.wraps.dify_config", SimpleNamespace(EDITION="CLOUD")),
patch("controllers.console.wraps.FeatureService.get_system_features", return_value=wraps_features),
):
with app.test_request_context(
"/forgot-password",
method="POST",
json={"email": "User@Example.com", "language": "zh-Hans"},
):
response = ForgotPasswordSendEmailApi().post()
assert response == {"result": "success", "data": "token-123"}
mock_get_account.assert_called_once_with("User@Example.com", session=mock_session)
mock_send_email.assert_called_once_with(
account=mock_account,
email="user@example.com",
language="zh-Hans",
is_allow_register=True,
)
mock_is_ip_limit.assert_called_once_with("127.0.0.1")
mock_extract_ip.assert_called_once()
class TestForgotPasswordCheckApi:
@patch("controllers.console.auth.forgot_password.AccountService.reset_forgot_password_error_rate_limit")
@patch("controllers.console.auth.forgot_password.AccountService.generate_reset_password_token")
@patch("controllers.console.auth.forgot_password.AccountService.revoke_reset_password_token")
@patch("controllers.console.auth.forgot_password.AccountService.add_forgot_password_error_rate_limit")
@patch("controllers.console.auth.forgot_password.AccountService.get_reset_password_data")
@patch("controllers.console.auth.forgot_password.AccountService.is_forgot_password_error_rate_limit")
def test_check_normalizes_email(
self,
mock_rate_limit_check,
mock_get_data,
mock_add_rate,
mock_revoke_token,
mock_generate_token,
mock_reset_rate,
app,
):
mock_rate_limit_check.return_value = False
mock_get_data.return_value = {"email": "Admin@Example.com", "code": "4321"}
mock_generate_token.return_value = (None, "new-token")
wraps_features = SimpleNamespace(enable_email_password_login=True)
with (
patch("controllers.console.wraps.dify_config", SimpleNamespace(EDITION="CLOUD")),
patch("controllers.console.wraps.FeatureService.get_system_features", return_value=wraps_features),
):
with app.test_request_context(
"/forgot-password/validity",
method="POST",
json={"email": "ADMIN@Example.com", "code": "4321", "token": "token-123"},
):
response = ForgotPasswordCheckApi().post()
assert response == {"is_valid": True, "email": "admin@example.com", "token": "new-token"}
mock_rate_limit_check.assert_called_once_with("admin@example.com")
mock_generate_token.assert_called_once_with(
"Admin@Example.com",
code="4321",
additional_data={"phase": "reset"},
)
mock_reset_rate.assert_called_once_with("admin@example.com")
mock_add_rate.assert_not_called()
mock_revoke_token.assert_called_once_with("token-123")
class TestForgotPasswordResetApi:
@patch("controllers.console.auth.forgot_password.ForgotPasswordResetApi._update_existing_account")
@patch("controllers.console.auth.forgot_password.Session")
@patch("controllers.console.auth.forgot_password.AccountService.get_account_by_email_with_case_fallback")
@patch("controllers.console.auth.forgot_password.AccountService.revoke_reset_password_token")
@patch("controllers.console.auth.forgot_password.AccountService.get_reset_password_data")
def test_reset_fetches_account_with_original_email(
self,
mock_get_reset_data,
mock_revoke_token,
mock_get_account,
mock_session_cls,
mock_update_account,
app,
):
mock_get_reset_data.return_value = {"phase": "reset", "email": "User@Example.com"}
mock_account = MagicMock()
mock_get_account.return_value = mock_account
mock_session = MagicMock()
mock_session_cls.return_value.__enter__.return_value = mock_session
wraps_features = SimpleNamespace(enable_email_password_login=True)
with (
patch("controllers.console.auth.forgot_password.db", SimpleNamespace(engine="engine")),
patch("controllers.console.wraps.dify_config", SimpleNamespace(EDITION="CLOUD")),
patch("controllers.console.wraps.FeatureService.get_system_features", return_value=wraps_features),
):
with app.test_request_context(
"/forgot-password/resets",
method="POST",
json={
"token": "token-123",
"new_password": "ValidPass123!",
"password_confirm": "ValidPass123!",
},
):
response = ForgotPasswordResetApi().post()
assert response == {"result": "success"}
mock_get_reset_data.assert_called_once_with("token-123")
mock_revoke_token.assert_called_once_with("token-123")
mock_get_account.assert_called_once_with("User@Example.com", session=mock_session)
mock_update_account.assert_called_once()
def test_get_account_by_email_with_case_fallback_uses_lowercase_lookup():
mock_session = MagicMock()
first_query = MagicMock()
first_query.scalar_one_or_none.return_value = None
expected_account = MagicMock()
second_query = MagicMock()
second_query.scalar_one_or_none.return_value = expected_account
mock_session.execute.side_effect = [first_query, second_query]
account = AccountService.get_account_by_email_with_case_fallback("Mixed@Test.com", session=mock_session)
assert account is expected_account
assert mock_session.execute.call_count == 2

View File

@ -76,7 +76,7 @@ class TestLoginApi:
@patch("controllers.console.wraps.db")
@patch("controllers.console.auth.login.dify_config.BILLING_ENABLED", False)
@patch("controllers.console.auth.login.AccountService.is_login_error_rate_limit")
@patch("controllers.console.auth.login.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.login.RegisterService.get_invitation_with_case_fallback")
@patch("controllers.console.auth.login.AccountService.authenticate")
@patch("controllers.console.auth.login.TenantService.get_join_tenants")
@patch("controllers.console.auth.login.AccountService.login")
@ -120,7 +120,7 @@ class TestLoginApi:
response = login_api.post()
# Assert
mock_authenticate.assert_called_once_with("test@example.com", "ValidPass123!")
mock_authenticate.assert_called_once_with("test@example.com", "ValidPass123!", None)
mock_login.assert_called_once()
mock_reset_rate_limit.assert_called_once_with("test@example.com")
assert response.json["result"] == "success"
@ -128,7 +128,7 @@ class TestLoginApi:
@patch("controllers.console.wraps.db")
@patch("controllers.console.auth.login.dify_config.BILLING_ENABLED", False)
@patch("controllers.console.auth.login.AccountService.is_login_error_rate_limit")
@patch("controllers.console.auth.login.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.login.RegisterService.get_invitation_with_case_fallback")
@patch("controllers.console.auth.login.AccountService.authenticate")
@patch("controllers.console.auth.login.TenantService.get_join_tenants")
@patch("controllers.console.auth.login.AccountService.login")
@ -182,7 +182,7 @@ class TestLoginApi:
@patch("controllers.console.wraps.db")
@patch("controllers.console.auth.login.dify_config.BILLING_ENABLED", False)
@patch("controllers.console.auth.login.AccountService.is_login_error_rate_limit")
@patch("controllers.console.auth.login.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.login.RegisterService.get_invitation_with_case_fallback")
def test_login_fails_when_rate_limited(self, mock_get_invitation, mock_is_rate_limit, mock_db, app):
"""
Test login rejection when rate limit is exceeded.
@ -230,7 +230,7 @@ class TestLoginApi:
@patch("controllers.console.wraps.db")
@patch("controllers.console.auth.login.dify_config.BILLING_ENABLED", False)
@patch("controllers.console.auth.login.AccountService.is_login_error_rate_limit")
@patch("controllers.console.auth.login.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.login.RegisterService.get_invitation_with_case_fallback")
@patch("controllers.console.auth.login.AccountService.authenticate")
@patch("controllers.console.auth.login.AccountService.add_login_error_rate_limit")
def test_login_fails_with_invalid_credentials(
@ -269,7 +269,7 @@ class TestLoginApi:
@patch("controllers.console.wraps.db")
@patch("controllers.console.auth.login.dify_config.BILLING_ENABLED", False)
@patch("controllers.console.auth.login.AccountService.is_login_error_rate_limit")
@patch("controllers.console.auth.login.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.login.RegisterService.get_invitation_with_case_fallback")
@patch("controllers.console.auth.login.AccountService.authenticate")
def test_login_fails_for_banned_account(
self, mock_authenticate, mock_get_invitation, mock_is_rate_limit, mock_db, app
@ -298,7 +298,7 @@ class TestLoginApi:
@patch("controllers.console.wraps.db")
@patch("controllers.console.auth.login.dify_config.BILLING_ENABLED", False)
@patch("controllers.console.auth.login.AccountService.is_login_error_rate_limit")
@patch("controllers.console.auth.login.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.login.RegisterService.get_invitation_with_case_fallback")
@patch("controllers.console.auth.login.AccountService.authenticate")
@patch("controllers.console.auth.login.TenantService.get_join_tenants")
@patch("controllers.console.auth.login.FeatureService.get_system_features")
@ -343,7 +343,7 @@ class TestLoginApi:
@patch("controllers.console.wraps.db")
@patch("controllers.console.auth.login.dify_config.BILLING_ENABLED", False)
@patch("controllers.console.auth.login.AccountService.is_login_error_rate_limit")
@patch("controllers.console.auth.login.RegisterService.get_invitation_if_token_valid")
@patch("controllers.console.auth.login.RegisterService.get_invitation_with_case_fallback")
def test_login_invitation_email_mismatch(self, mock_get_invitation, mock_is_rate_limit, mock_db, app):
"""
Test login failure when invitation email doesn't match login email.
@ -371,6 +371,52 @@ class TestLoginApi:
with pytest.raises(InvalidEmailError):
login_api.post()
@patch("controllers.console.wraps.db")
@patch("controllers.console.auth.login.dify_config.BILLING_ENABLED", False)
@patch("controllers.console.auth.login.AccountService.is_login_error_rate_limit")
@patch("controllers.console.auth.login.RegisterService.get_invitation_with_case_fallback")
@patch("controllers.console.auth.login.AccountService.authenticate")
@patch("controllers.console.auth.login.AccountService.add_login_error_rate_limit")
@patch("controllers.console.auth.login.TenantService.get_join_tenants")
@patch("controllers.console.auth.login.AccountService.login")
@patch("controllers.console.auth.login.AccountService.reset_login_error_rate_limit")
def test_login_retries_with_lowercase_email(
self,
mock_reset_rate_limit,
mock_login_service,
mock_get_tenants,
mock_add_rate_limit,
mock_authenticate,
mock_get_invitation,
mock_is_rate_limit,
mock_db,
app,
mock_account,
mock_token_pair,
):
"""Test that login retries with lowercase email when uppercase lookup fails."""
mock_db.session.query.return_value.first.return_value = MagicMock()
mock_is_rate_limit.return_value = False
mock_get_invitation.return_value = None
mock_authenticate.side_effect = [AccountPasswordError("Invalid"), mock_account]
mock_get_tenants.return_value = [MagicMock()]
mock_login_service.return_value = mock_token_pair
with app.test_request_context(
"/login",
method="POST",
json={"email": "Upper@Example.com", "password": encode_password("ValidPass123!")},
):
response = LoginApi().post()
assert response.json["result"] == "success"
assert mock_authenticate.call_args_list == [
(("Upper@Example.com", "ValidPass123!", None), {}),
(("upper@example.com", "ValidPass123!", None), {}),
]
mock_add_rate_limit.assert_not_called()
mock_reset_rate_limit.assert_called_once_with("upper@example.com")
class TestLogoutApi:
"""Test cases for the LogoutApi endpoint."""

View File

@ -12,6 +12,7 @@ from controllers.console.auth.oauth import (
)
from libs.oauth import OAuthUserInfo
from models.account import AccountStatus
from services.account_service import AccountService
from services.errors.account import AccountRegisterError
@ -171,7 +172,7 @@ class TestOAuthCallback:
):
mock_config.CONSOLE_WEB_URL = "http://localhost:3000"
mock_get_providers.return_value = {"github": oauth_setup["provider"]}
mock_generate_account.return_value = oauth_setup["account"]
mock_generate_account.return_value = (oauth_setup["account"], True)
mock_account_service.login.return_value = oauth_setup["token_pair"]
with app.test_request_context("/auth/oauth/github/callback?code=test_code"):
@ -179,7 +180,7 @@ class TestOAuthCallback:
oauth_setup["provider"].get_access_token.assert_called_once_with("test_code")
oauth_setup["provider"].get_user_info.assert_called_once_with("access_token")
mock_redirect.assert_called_once_with("http://localhost:3000")
mock_redirect.assert_called_once_with("http://localhost:3000?oauth_new_user=true")
@pytest.mark.parametrize(
("exception", "expected_error"),
@ -215,6 +216,34 @@ class TestOAuthCallback:
assert status_code == 400
assert response["error"] == expected_error
@patch("controllers.console.auth.oauth.dify_config")
@patch("controllers.console.auth.oauth.get_oauth_providers")
@patch("controllers.console.auth.oauth.RegisterService")
@patch("controllers.console.auth.oauth.redirect")
def test_invitation_comparison_is_case_insensitive(
self,
mock_redirect,
mock_register_service,
mock_get_providers,
mock_config,
resource,
app,
oauth_setup,
):
mock_config.CONSOLE_WEB_URL = "http://localhost:3000"
oauth_setup["provider"].get_user_info.return_value = OAuthUserInfo(
id="123", name="Test User", email="User@Example.com"
)
mock_get_providers.return_value = {"github": oauth_setup["provider"]}
mock_register_service.is_valid_invite_token.return_value = True
mock_register_service.get_invitation_by_token.return_value = {"email": "user@example.com"}
with app.test_request_context("/auth/oauth/github/callback?code=test_code&state=invite123"):
resource.get("github")
mock_register_service.get_invitation_by_token.assert_called_once_with(token="invite123")
mock_redirect.assert_called_once_with("http://localhost:3000/signin/invite-settings?invite_token=invite123")
@pytest.mark.parametrize(
("account_status", "expected_redirect"),
[
@ -223,7 +252,7 @@ class TestOAuthCallback:
# This documents actual behavior. See test_defensive_check_for_closed_account_status for details
(
AccountStatus.CLOSED.value,
"http://localhost:3000",
"http://localhost:3000?oauth_new_user=false",
),
],
)
@ -260,7 +289,7 @@ class TestOAuthCallback:
account = MagicMock()
account.status = account_status
account.id = "123"
mock_generate_account.return_value = account
mock_generate_account.return_value = (account, False)
# Mock login for CLOSED status
mock_token_pair = MagicMock()
@ -296,7 +325,7 @@ class TestOAuthCallback:
mock_account = MagicMock()
mock_account.status = AccountStatus.PENDING
mock_generate_account.return_value = mock_account
mock_generate_account.return_value = (mock_account, False)
mock_token_pair = MagicMock()
mock_token_pair.access_token = "jwt_access_token"
@ -360,7 +389,7 @@ class TestOAuthCallback:
closed_account.status = AccountStatus.CLOSED
closed_account.id = "123"
closed_account.name = "Closed Account"
mock_generate_account.return_value = closed_account
mock_generate_account.return_value = (closed_account, False)
# Mock successful login (current behavior)
mock_token_pair = MagicMock()
@ -374,7 +403,7 @@ class TestOAuthCallback:
resource.get("github")
# Verify current behavior: login succeeds (this is NOT ideal)
mock_redirect.assert_called_once_with("http://localhost:3000")
mock_redirect.assert_called_once_with("http://localhost:3000?oauth_new_user=false")
mock_account_service.login.assert_called_once()
# Document expected behavior in comments:
@ -395,12 +424,12 @@ class TestAccountGeneration:
account.name = "Test User"
return account
@patch("controllers.console.auth.oauth.db")
@patch("controllers.console.auth.oauth.Account")
@patch("controllers.console.auth.oauth.AccountService.get_account_by_email_with_case_fallback")
@patch("controllers.console.auth.oauth.Session")
@patch("controllers.console.auth.oauth.select")
@patch("controllers.console.auth.oauth.Account")
@patch("controllers.console.auth.oauth.db")
def test_should_get_account_by_openid_or_email(
self, mock_select, mock_session, mock_account_model, mock_db, user_info, mock_account
self, mock_db, mock_account_model, mock_session, mock_get_account, user_info, mock_account
):
# Mock db.engine for Session creation
mock_db.engine = MagicMock()
@ -410,15 +439,31 @@ class TestAccountGeneration:
result = _get_account_by_openid_or_email("github", user_info)
assert result == mock_account
mock_account_model.get_by_openid.assert_called_once_with("github", "123")
mock_get_account.assert_not_called()
# Test fallback to email
# Test fallback to email lookup
mock_account_model.get_by_openid.return_value = None
mock_session_instance = MagicMock()
mock_session_instance.execute.return_value.scalar_one_or_none.return_value = mock_account
mock_session.return_value.__enter__.return_value = mock_session_instance
mock_get_account.return_value = mock_account
result = _get_account_by_openid_or_email("github", user_info)
assert result == mock_account
mock_get_account.assert_called_once_with(user_info.email, session=mock_session_instance)
def test_get_account_by_email_with_case_fallback_uses_lowercase_lookup(self):
mock_session = MagicMock()
first_result = MagicMock()
first_result.scalar_one_or_none.return_value = None
expected_account = MagicMock()
second_result = MagicMock()
second_result.scalar_one_or_none.return_value = expected_account
mock_session.execute.side_effect = [first_result, second_result]
result = AccountService.get_account_by_email_with_case_fallback("Case@Test.com", session=mock_session)
assert result == expected_account
assert mock_session.execute.call_count == 2
@pytest.mark.parametrize(
("allow_register", "existing_account", "should_create"),
@ -458,13 +503,43 @@ class TestAccountGeneration:
with pytest.raises(AccountRegisterError):
_generate_account("github", user_info)
else:
result = _generate_account("github", user_info)
result, oauth_new_user = _generate_account("github", user_info)
assert result == mock_account
assert oauth_new_user == should_create
if should_create:
mock_register_service.register.assert_called_once_with(
email="test@example.com", name="Test User", password=None, open_id="123", provider="github"
)
else:
mock_register_service.register.assert_not_called()
@patch("controllers.console.auth.oauth._get_account_by_openid_or_email", return_value=None)
@patch("controllers.console.auth.oauth.FeatureService")
@patch("controllers.console.auth.oauth.RegisterService")
@patch("controllers.console.auth.oauth.AccountService")
@patch("controllers.console.auth.oauth.TenantService")
@patch("controllers.console.auth.oauth.db")
def test_should_register_with_lowercase_email(
self,
mock_db,
mock_tenant_service,
mock_account_service,
mock_register_service,
mock_feature_service,
mock_get_account,
app,
):
user_info = OAuthUserInfo(id="123", name="Test User", email="Upper@Example.com")
mock_feature_service.get_system_features.return_value.is_allow_register = True
mock_register_service.register.return_value = MagicMock()
with app.test_request_context(headers={"Accept-Language": "en-US"}):
_generate_account("github", user_info)
mock_register_service.register.assert_called_once_with(
email="upper@example.com", name="Test User", password=None, open_id="123", provider="github"
)
@patch("controllers.console.auth.oauth._get_account_by_openid_or_email")
@patch("controllers.console.auth.oauth.TenantService")
@ -490,9 +565,10 @@ class TestAccountGeneration:
mock_tenant_service.create_tenant.return_value = mock_new_tenant
with app.test_request_context(headers={"Accept-Language": "en-US,en;q=0.9"}):
result = _generate_account("github", user_info)
result, oauth_new_user = _generate_account("github", user_info)
assert result == mock_account
assert oauth_new_user is False
mock_tenant_service.create_tenant.assert_called_once_with("Test User's Workspace")
mock_tenant_service.create_tenant_member.assert_called_once_with(
mock_new_tenant, mock_account, role="owner"

View File

@ -28,6 +28,22 @@ from controllers.console.auth.forgot_password import (
from controllers.console.error import AccountNotFound, EmailSendIpLimitError
@pytest.fixture(autouse=True)
def _mock_forgot_password_session():
with patch("controllers.console.auth.forgot_password.Session") as mock_session_cls:
mock_session = MagicMock()
mock_session_cls.return_value.__enter__.return_value = mock_session
mock_session_cls.return_value.__exit__.return_value = None
yield mock_session
@pytest.fixture(autouse=True)
def _mock_forgot_password_db():
with patch("controllers.console.auth.forgot_password.db") as mock_db:
mock_db.engine = MagicMock()
yield mock_db
class TestForgotPasswordSendEmailApi:
"""Test cases for sending password reset emails."""
@ -47,20 +63,16 @@ class TestForgotPasswordSendEmailApi:
return account
@patch("controllers.console.wraps.db")
@patch("controllers.console.auth.forgot_password.db")
@patch("controllers.console.auth.forgot_password.AccountService.is_email_send_ip_limit")
@patch("controllers.console.auth.forgot_password.Session")
@patch("controllers.console.auth.forgot_password.select")
@patch("controllers.console.auth.forgot_password.AccountService.get_account_by_email_with_case_fallback")
@patch("controllers.console.auth.forgot_password.AccountService.send_reset_password_email")
@patch("controllers.console.auth.forgot_password.FeatureService.get_system_features")
def test_send_reset_email_success(
self,
mock_get_features,
mock_send_email,
mock_select,
mock_session,
mock_get_account,
mock_is_ip_limit,
mock_forgot_db,
mock_wraps_db,
app,
mock_account,
@ -75,11 +87,8 @@ class TestForgotPasswordSendEmailApi:
"""
# Arrange
mock_wraps_db.session.query.return_value.first.return_value = MagicMock()
mock_forgot_db.engine = MagicMock()
mock_is_ip_limit.return_value = False
mock_session_instance = MagicMock()
mock_session_instance.execute.return_value.scalar_one_or_none.return_value = mock_account
mock_session.return_value.__enter__.return_value = mock_session_instance
mock_get_account.return_value = mock_account
mock_send_email.return_value = "reset_token_123"
mock_get_features.return_value.is_allow_register = True
@ -125,20 +134,16 @@ class TestForgotPasswordSendEmailApi:
],
)
@patch("controllers.console.wraps.db")
@patch("controllers.console.auth.forgot_password.db")
@patch("controllers.console.auth.forgot_password.AccountService.is_email_send_ip_limit")
@patch("controllers.console.auth.forgot_password.Session")
@patch("controllers.console.auth.forgot_password.select")
@patch("controllers.console.auth.forgot_password.AccountService.get_account_by_email_with_case_fallback")
@patch("controllers.console.auth.forgot_password.AccountService.send_reset_password_email")
@patch("controllers.console.auth.forgot_password.FeatureService.get_system_features")
def test_send_reset_email_language_handling(
self,
mock_get_features,
mock_send_email,
mock_select,
mock_session,
mock_get_account,
mock_is_ip_limit,
mock_forgot_db,
mock_wraps_db,
app,
mock_account,
@ -154,11 +159,8 @@ class TestForgotPasswordSendEmailApi:
"""
# Arrange
mock_wraps_db.session.query.return_value.first.return_value = MagicMock()
mock_forgot_db.engine = MagicMock()
mock_is_ip_limit.return_value = False
mock_session_instance = MagicMock()
mock_session_instance.execute.return_value.scalar_one_or_none.return_value = mock_account
mock_session.return_value.__enter__.return_value = mock_session_instance
mock_get_account.return_value = mock_account
mock_send_email.return_value = "token"
mock_get_features.return_value.is_allow_register = True
@ -229,8 +231,46 @@ class TestForgotPasswordCheckApi:
assert response["email"] == "test@example.com"
assert response["token"] == "new_token"
mock_revoke_token.assert_called_once_with("old_token")
mock_generate_token.assert_called_once_with(
"test@example.com", code="123456", additional_data={"phase": "reset"}
)
mock_reset_rate_limit.assert_called_once_with("test@example.com")
@patch("controllers.console.wraps.db")
@patch("controllers.console.auth.forgot_password.AccountService.is_forgot_password_error_rate_limit")
@patch("controllers.console.auth.forgot_password.AccountService.get_reset_password_data")
@patch("controllers.console.auth.forgot_password.AccountService.revoke_reset_password_token")
@patch("controllers.console.auth.forgot_password.AccountService.generate_reset_password_token")
@patch("controllers.console.auth.forgot_password.AccountService.reset_forgot_password_error_rate_limit")
def test_verify_code_preserves_token_email_case(
self,
mock_reset_rate_limit,
mock_generate_token,
mock_revoke_token,
mock_get_data,
mock_is_rate_limit,
mock_db,
app,
):
mock_db.session.query.return_value.first.return_value = MagicMock()
mock_is_rate_limit.return_value = False
mock_get_data.return_value = {"email": "User@Example.com", "code": "999888"}
mock_generate_token.return_value = (None, "fresh-token")
with app.test_request_context(
"/forgot-password/validity",
method="POST",
json={"email": "user@example.com", "code": "999888", "token": "upper_token"},
):
response = ForgotPasswordCheckApi().post()
assert response == {"is_valid": True, "email": "user@example.com", "token": "fresh-token"}
mock_generate_token.assert_called_once_with(
"User@Example.com", code="999888", additional_data={"phase": "reset"}
)
mock_revoke_token.assert_called_once_with("upper_token")
mock_reset_rate_limit.assert_called_once_with("user@example.com")
@patch("controllers.console.wraps.db")
@patch("controllers.console.auth.forgot_password.AccountService.is_forgot_password_error_rate_limit")
def test_verify_code_rate_limited(self, mock_is_rate_limit, mock_db, app):
@ -355,20 +395,16 @@ class TestForgotPasswordResetApi:
return account
@patch("controllers.console.wraps.db")
@patch("controllers.console.auth.forgot_password.db")
@patch("controllers.console.auth.forgot_password.AccountService.get_reset_password_data")
@patch("controllers.console.auth.forgot_password.AccountService.revoke_reset_password_token")
@patch("controllers.console.auth.forgot_password.Session")
@patch("controllers.console.auth.forgot_password.select")
@patch("controllers.console.auth.forgot_password.AccountService.get_account_by_email_with_case_fallback")
@patch("controllers.console.auth.forgot_password.TenantService.get_join_tenants")
def test_reset_password_success(
self,
mock_get_tenants,
mock_select,
mock_session,
mock_get_account,
mock_revoke_token,
mock_get_data,
mock_forgot_db,
mock_wraps_db,
app,
mock_account,
@ -383,11 +419,8 @@ class TestForgotPasswordResetApi:
"""
# Arrange
mock_wraps_db.session.query.return_value.first.return_value = MagicMock()
mock_forgot_db.engine = MagicMock()
mock_get_data.return_value = {"email": "test@example.com", "phase": "reset"}
mock_session_instance = MagicMock()
mock_session_instance.execute.return_value.scalar_one_or_none.return_value = mock_account
mock_session.return_value.__enter__.return_value = mock_session_instance
mock_get_account.return_value = mock_account
mock_get_tenants.return_value = [MagicMock()]
# Act
@ -475,13 +508,11 @@ class TestForgotPasswordResetApi:
api.post()
@patch("controllers.console.wraps.db")
@patch("controllers.console.auth.forgot_password.db")
@patch("controllers.console.auth.forgot_password.AccountService.get_reset_password_data")
@patch("controllers.console.auth.forgot_password.AccountService.revoke_reset_password_token")
@patch("controllers.console.auth.forgot_password.Session")
@patch("controllers.console.auth.forgot_password.select")
@patch("controllers.console.auth.forgot_password.AccountService.get_account_by_email_with_case_fallback")
def test_reset_password_account_not_found(
self, mock_select, mock_session, mock_revoke_token, mock_get_data, mock_forgot_db, mock_wraps_db, app
self, mock_get_account, mock_revoke_token, mock_get_data, mock_wraps_db, app
):
"""
Test password reset for non-existent account.
@ -491,11 +522,8 @@ class TestForgotPasswordResetApi:
"""
# Arrange
mock_wraps_db.session.query.return_value.first.return_value = MagicMock()
mock_forgot_db.engine = MagicMock()
mock_get_data.return_value = {"email": "nonexistent@example.com", "phase": "reset"}
mock_session_instance = MagicMock()
mock_session_instance.execute.return_value.scalar_one_or_none.return_value = None
mock_session.return_value.__enter__.return_value = mock_session_instance
mock_get_account.return_value = None
# Act & Assert
with app.test_request_context(

View File

@ -0,0 +1 @@
"""Unit tests for `controllers.console.datasets` controllers."""

View File

@ -0,0 +1,430 @@
"""
Unit tests for the dataset document download endpoint.
These tests validate that the controller returns a signed download URL for
upload-file documents, and rejects unsupported or missing file cases.
"""
from __future__ import annotations
import importlib
import sys
from collections import UserDict
from io import BytesIO
from types import SimpleNamespace
from typing import Any
from zipfile import ZipFile
import pytest
from flask import Flask
from werkzeug.exceptions import Forbidden, NotFound
@pytest.fixture
def app() -> Flask:
"""Create a minimal Flask app for request-context based controller tests."""
app = Flask(__name__)
app.config["TESTING"] = True
return app
@pytest.fixture
def datasets_document_module(monkeypatch: pytest.MonkeyPatch):
"""
Reload `controllers.console.datasets.datasets_document` with lightweight decorators.
We patch auth / setup / rate-limit decorators to no-ops so we can unit test the
controller logic without requiring the full console stack.
"""
from controllers.console import console_ns, wraps
from libs import login
def _noop(func): # type: ignore[no-untyped-def]
return func
# Bypass login/setup/account checks in unit tests.
monkeypatch.setattr(login, "login_required", _noop)
monkeypatch.setattr(wraps, "setup_required", _noop)
monkeypatch.setattr(wraps, "account_initialization_required", _noop)
# Bypass billing-related decorators used by other endpoints in this module.
monkeypatch.setattr(wraps, "cloud_edition_billing_resource_check", lambda *_args, **_kwargs: (lambda f: f))
monkeypatch.setattr(wraps, "cloud_edition_billing_rate_limit_check", lambda *_args, **_kwargs: (lambda f: f))
# Avoid Flask-RESTX route registration side effects during import.
def _noop_route(*_args, **_kwargs): # type: ignore[override]
def _decorator(cls):
return cls
return _decorator
monkeypatch.setattr(console_ns, "route", _noop_route)
module_name = "controllers.console.datasets.datasets_document"
sys.modules.pop(module_name, None)
return importlib.import_module(module_name)
def _mock_user(*, is_dataset_editor: bool = True) -> SimpleNamespace:
"""Build a minimal user object compatible with dataset permission checks."""
return SimpleNamespace(is_dataset_editor=is_dataset_editor, id="user-123")
def _mock_document(
*,
document_id: str,
tenant_id: str,
data_source_type: str,
upload_file_id: str | None,
) -> SimpleNamespace:
"""Build a minimal document object used by the controller."""
data_source_info_dict: dict[str, Any] | None = None
if upload_file_id is not None:
data_source_info_dict = {"upload_file_id": upload_file_id}
else:
data_source_info_dict = {}
return SimpleNamespace(
id=document_id,
tenant_id=tenant_id,
data_source_type=data_source_type,
data_source_info_dict=data_source_info_dict,
)
def _wire_common_success_mocks(
*,
module,
monkeypatch: pytest.MonkeyPatch,
current_tenant_id: str,
document_tenant_id: str,
data_source_type: str,
upload_file_id: str | None,
upload_file_exists: bool,
signed_url: str,
) -> None:
"""Patch controller dependencies to create a deterministic test environment."""
import services.dataset_service as dataset_service_module
# Make `current_account_with_tenant()` return a known user + tenant id.
monkeypatch.setattr(module, "current_account_with_tenant", lambda: (_mock_user(), current_tenant_id))
# Return a dataset object and allow permission checks to pass.
monkeypatch.setattr(module.DatasetService, "get_dataset", lambda _dataset_id: SimpleNamespace(id="ds-1"))
monkeypatch.setattr(module.DatasetService, "check_dataset_permission", lambda *_args, **_kwargs: None)
# Return a document that will be validated inside DocumentResource.get_document.
document = _mock_document(
document_id="doc-1",
tenant_id=document_tenant_id,
data_source_type=data_source_type,
upload_file_id=upload_file_id,
)
monkeypatch.setattr(module.DocumentService, "get_document", lambda *_args, **_kwargs: document)
# Mock UploadFile lookup via FileService batch helper.
upload_files_by_id: dict[str, Any] = {}
if upload_file_exists and upload_file_id is not None:
upload_files_by_id[str(upload_file_id)] = SimpleNamespace(id=str(upload_file_id))
monkeypatch.setattr(module.FileService, "get_upload_files_by_ids", lambda *_args, **_kwargs: upload_files_by_id)
# Mock signing helper so the returned URL is deterministic.
monkeypatch.setattr(dataset_service_module.file_helpers, "get_signed_file_url", lambda **_kwargs: signed_url)
def _mock_send_file(obj, **kwargs): # type: ignore[no-untyped-def]
"""Return a lightweight representation of `send_file(...)` for unit tests."""
class _ResponseMock(UserDict):
def __init__(self, sent_file: object, send_file_kwargs: dict[str, object]) -> None:
super().__init__({"_sent_file": sent_file, "_send_file_kwargs": send_file_kwargs})
self._on_close: object | None = None
def call_on_close(self, func): # type: ignore[no-untyped-def]
self._on_close = func
return func
return _ResponseMock(obj, kwargs)
def test_batch_download_zip_returns_send_file(
app: Flask, datasets_document_module, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Ensure batch ZIP download returns a zip attachment via `send_file`."""
# Arrange common permission mocks.
monkeypatch.setattr(datasets_document_module, "current_account_with_tenant", lambda: (_mock_user(), "tenant-123"))
monkeypatch.setattr(
datasets_document_module.DatasetService, "get_dataset", lambda _dataset_id: SimpleNamespace(id="ds-1")
)
monkeypatch.setattr(
datasets_document_module.DatasetService, "check_dataset_permission", lambda *_args, **_kwargs: None
)
# Two upload-file documents, each referencing an UploadFile.
doc1 = _mock_document(
document_id="11111111-1111-1111-1111-111111111111",
tenant_id="tenant-123",
data_source_type="upload_file",
upload_file_id="file-1",
)
doc2 = _mock_document(
document_id="22222222-2222-2222-2222-222222222222",
tenant_id="tenant-123",
data_source_type="upload_file",
upload_file_id="file-2",
)
monkeypatch.setattr(
datasets_document_module.DocumentService,
"get_documents_by_ids",
lambda *_args, **_kwargs: [doc1, doc2],
)
monkeypatch.setattr(
datasets_document_module.FileService,
"get_upload_files_by_ids",
lambda *_args, **_kwargs: {
"file-1": SimpleNamespace(id="file-1", name="a.txt", key="k1"),
"file-2": SimpleNamespace(id="file-2", name="b.txt", key="k2"),
},
)
# Mock storage streaming content.
import services.file_service as file_service_module
monkeypatch.setattr(file_service_module.storage, "load", lambda _key, stream=True: [b"hello"])
# Replace send_file used by the controller to avoid a real Flask response object.
monkeypatch.setattr(datasets_document_module, "send_file", _mock_send_file)
# Act
with app.test_request_context(
"/datasets/ds-1/documents/download-zip",
method="POST",
json={"document_ids": ["11111111-1111-1111-1111-111111111111", "22222222-2222-2222-2222-222222222222"]},
):
api = datasets_document_module.DocumentBatchDownloadZipApi()
result = api.post(dataset_id="ds-1")
# Assert: we returned via send_file with correct mime type and attachment.
assert result["_send_file_kwargs"]["mimetype"] == "application/zip"
assert result["_send_file_kwargs"]["as_attachment"] is True
assert isinstance(result["_send_file_kwargs"]["download_name"], str)
assert result["_send_file_kwargs"]["download_name"].endswith(".zip")
# Ensure our cleanup hook is registered and execute it to avoid temp file leaks in unit tests.
assert getattr(result, "_on_close", None) is not None
result._on_close() # type: ignore[attr-defined]
def test_batch_download_zip_response_is_openable_zip(
app: Flask, datasets_document_module, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Ensure the real Flask `send_file` response body is a valid ZIP that can be opened."""
# Arrange: same controller mocks as the lightweight send_file test, but we keep the real `send_file`.
monkeypatch.setattr(datasets_document_module, "current_account_with_tenant", lambda: (_mock_user(), "tenant-123"))
monkeypatch.setattr(
datasets_document_module.DatasetService, "get_dataset", lambda _dataset_id: SimpleNamespace(id="ds-1")
)
monkeypatch.setattr(
datasets_document_module.DatasetService, "check_dataset_permission", lambda *_args, **_kwargs: None
)
doc1 = _mock_document(
document_id="33333333-3333-3333-3333-333333333333",
tenant_id="tenant-123",
data_source_type="upload_file",
upload_file_id="file-1",
)
doc2 = _mock_document(
document_id="44444444-4444-4444-4444-444444444444",
tenant_id="tenant-123",
data_source_type="upload_file",
upload_file_id="file-2",
)
monkeypatch.setattr(
datasets_document_module.DocumentService,
"get_documents_by_ids",
lambda *_args, **_kwargs: [doc1, doc2],
)
monkeypatch.setattr(
datasets_document_module.FileService,
"get_upload_files_by_ids",
lambda *_args, **_kwargs: {
"file-1": SimpleNamespace(id="file-1", name="a.txt", key="k1"),
"file-2": SimpleNamespace(id="file-2", name="b.txt", key="k2"),
},
)
# Stream distinct bytes per key so we can verify both ZIP entries.
import services.file_service as file_service_module
monkeypatch.setattr(
file_service_module.storage, "load", lambda key, stream=True: [b"one"] if key == "k1" else [b"two"]
)
# Act
with app.test_request_context(
"/datasets/ds-1/documents/download-zip",
method="POST",
json={"document_ids": ["33333333-3333-3333-3333-333333333333", "44444444-4444-4444-4444-444444444444"]},
):
api = datasets_document_module.DocumentBatchDownloadZipApi()
response = api.post(dataset_id="ds-1")
# Assert: response body is a valid ZIP and contains the expected entries.
response.direct_passthrough = False
data = response.get_data()
response.close()
with ZipFile(BytesIO(data), mode="r") as zf:
assert zf.namelist() == ["a.txt", "b.txt"]
assert zf.read("a.txt") == b"one"
assert zf.read("b.txt") == b"two"
def test_batch_download_zip_rejects_non_upload_file_document(
app: Flask, datasets_document_module, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Ensure batch ZIP download rejects non upload-file documents."""
monkeypatch.setattr(datasets_document_module, "current_account_with_tenant", lambda: (_mock_user(), "tenant-123"))
monkeypatch.setattr(
datasets_document_module.DatasetService, "get_dataset", lambda _dataset_id: SimpleNamespace(id="ds-1")
)
monkeypatch.setattr(
datasets_document_module.DatasetService, "check_dataset_permission", lambda *_args, **_kwargs: None
)
doc = _mock_document(
document_id="55555555-5555-5555-5555-555555555555",
tenant_id="tenant-123",
data_source_type="website_crawl",
upload_file_id="file-1",
)
monkeypatch.setattr(
datasets_document_module.DocumentService,
"get_documents_by_ids",
lambda *_args, **_kwargs: [doc],
)
with app.test_request_context(
"/datasets/ds-1/documents/download-zip",
method="POST",
json={"document_ids": ["55555555-5555-5555-5555-555555555555"]},
):
api = datasets_document_module.DocumentBatchDownloadZipApi()
with pytest.raises(NotFound):
api.post(dataset_id="ds-1")
def test_document_download_returns_url_for_upload_file_document(
app: Flask, datasets_document_module, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Ensure upload-file documents return a `{url}` JSON payload."""
_wire_common_success_mocks(
module=datasets_document_module,
monkeypatch=monkeypatch,
current_tenant_id="tenant-123",
document_tenant_id="tenant-123",
data_source_type="upload_file",
upload_file_id="file-123",
upload_file_exists=True,
signed_url="https://example.com/signed",
)
# Build a request context then call the resource method directly.
with app.test_request_context("/datasets/ds-1/documents/doc-1/download", method="GET"):
api = datasets_document_module.DocumentDownloadApi()
result = api.get(dataset_id="ds-1", document_id="doc-1")
assert result == {"url": "https://example.com/signed"}
def test_document_download_rejects_non_upload_file_document(
app: Flask, datasets_document_module, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Ensure non-upload documents raise 404 (no file to download)."""
_wire_common_success_mocks(
module=datasets_document_module,
monkeypatch=monkeypatch,
current_tenant_id="tenant-123",
document_tenant_id="tenant-123",
data_source_type="website_crawl",
upload_file_id="file-123",
upload_file_exists=True,
signed_url="https://example.com/signed",
)
with app.test_request_context("/datasets/ds-1/documents/doc-1/download", method="GET"):
api = datasets_document_module.DocumentDownloadApi()
with pytest.raises(NotFound):
api.get(dataset_id="ds-1", document_id="doc-1")
def test_document_download_rejects_missing_upload_file_id(
app: Flask, datasets_document_module, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Ensure missing `upload_file_id` raises 404."""
_wire_common_success_mocks(
module=datasets_document_module,
monkeypatch=monkeypatch,
current_tenant_id="tenant-123",
document_tenant_id="tenant-123",
data_source_type="upload_file",
upload_file_id=None,
upload_file_exists=False,
signed_url="https://example.com/signed",
)
with app.test_request_context("/datasets/ds-1/documents/doc-1/download", method="GET"):
api = datasets_document_module.DocumentDownloadApi()
with pytest.raises(NotFound):
api.get(dataset_id="ds-1", document_id="doc-1")
def test_document_download_rejects_when_upload_file_record_missing(
app: Flask, datasets_document_module, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Ensure missing UploadFile row raises 404."""
_wire_common_success_mocks(
module=datasets_document_module,
monkeypatch=monkeypatch,
current_tenant_id="tenant-123",
document_tenant_id="tenant-123",
data_source_type="upload_file",
upload_file_id="file-123",
upload_file_exists=False,
signed_url="https://example.com/signed",
)
with app.test_request_context("/datasets/ds-1/documents/doc-1/download", method="GET"):
api = datasets_document_module.DocumentDownloadApi()
with pytest.raises(NotFound):
api.get(dataset_id="ds-1", document_id="doc-1")
def test_document_download_rejects_tenant_mismatch(
app: Flask, datasets_document_module, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Ensure tenant mismatch is rejected by the shared `get_document()` permission check."""
_wire_common_success_mocks(
module=datasets_document_module,
monkeypatch=monkeypatch,
current_tenant_id="tenant-123",
document_tenant_id="tenant-999",
data_source_type="upload_file",
upload_file_id="file-123",
upload_file_exists=True,
signed_url="https://example.com/signed",
)
with app.test_request_context("/datasets/ds-1/documents/doc-1/download", method="GET"):
api = datasets_document_module.DocumentDownloadApi()
with pytest.raises(Forbidden):
api.get(dataset_id="ds-1", document_id="doc-1")

View File

@ -0,0 +1,49 @@
from __future__ import annotations
"""
Unit tests for the external dataset controller payload schemas.
These tests focus on Pydantic validation rules so we can catch regressions
in request constraints (e.g. max length changes) without exercising the
full Flask/RESTX request stack.
"""
import pytest
from pydantic import ValidationError
from controllers.console.datasets.external import ExternalDatasetCreatePayload
def test_external_dataset_create_payload_allows_name_length_100() -> None:
"""Ensure the `name` field accepts up to 100 characters (inclusive)."""
# Build a request payload with a boundary-length name value.
name_100: str = "a" * 100
payload = {
"external_knowledge_api_id": "ek-api-1",
"external_knowledge_id": "ek-1",
"name": name_100,
}
model = ExternalDatasetCreatePayload.model_validate(payload)
assert model.name == name_100
def test_external_dataset_create_payload_rejects_name_length_101() -> None:
"""Ensure the `name` field rejects values longer than 100 characters."""
# Build a request payload that exceeds the max length by 1.
name_101: str = "a" * 101
payload: dict[str, object] = {
"external_knowledge_api_id": "ek-api-1",
"external_knowledge_id": "ek-1",
"name": name_101,
}
with pytest.raises(ValidationError) as exc_info:
ExternalDatasetCreatePayload.model_validate(payload)
errors = exc_info.value.errors()
assert errors[0]["loc"] == ("name",)
assert errors[0]["type"] == "string_too_long"
assert errors[0]["ctx"]["max_length"] == 100

View File

@ -0,0 +1,145 @@
"""
Test for document detail API data_source_info serialization fix.
This test verifies that the document detail API returns both data_source_info
and data_source_detail_dict for all data_source_type values, including "local_file".
"""
import json
from typing import Generic, Literal, NotRequired, TypedDict, TypeVar, Union
from models.dataset import Document
class LocalFileInfo(TypedDict):
file_path: str
size: int
created_at: NotRequired[str]
class UploadFileInfo(TypedDict):
upload_file_id: str
class NotionImportInfo(TypedDict):
notion_page_id: str
workspace_id: str
class WebsiteCrawlInfo(TypedDict):
url: str
job_id: str
RawInfo = Union[LocalFileInfo, UploadFileInfo, NotionImportInfo, WebsiteCrawlInfo]
T_type = TypeVar("T_type", bound=str)
T_info = TypeVar("T_info", bound=Union[LocalFileInfo, UploadFileInfo, NotionImportInfo, WebsiteCrawlInfo])
class Case(TypedDict, Generic[T_type, T_info]):
data_source_type: T_type
data_source_info: str
expected_raw: T_info
LocalFileCase = Case[Literal["local_file"], LocalFileInfo]
UploadFileCase = Case[Literal["upload_file"], UploadFileInfo]
NotionImportCase = Case[Literal["notion_import"], NotionImportInfo]
WebsiteCrawlCase = Case[Literal["website_crawl"], WebsiteCrawlInfo]
AnyCase = Union[LocalFileCase, UploadFileCase, NotionImportCase, WebsiteCrawlCase]
case_1: LocalFileCase = {
"data_source_type": "local_file",
"data_source_info": json.dumps({"file_path": "/tmp/test.txt", "size": 1024}),
"expected_raw": {"file_path": "/tmp/test.txt", "size": 1024},
}
# ERROR: Expected LocalFileInfo, but got WebsiteCrawlInfo
case_2: LocalFileCase = {
"data_source_type": "local_file",
"data_source_info": "...",
"expected_raw": {"file_path": "https://google.com", "size": 123},
}
cases: list[AnyCase] = [case_1]
class TestDocumentDetailDataSourceInfo:
"""Test cases for document detail API data_source_info serialization."""
def test_data_source_info_dict_returns_raw_data(self):
"""Test that data_source_info_dict returns raw JSON data for all data_source_type values."""
# Test data for different data_source_type values
for case in cases:
document = Document(
data_source_type=case["data_source_type"],
data_source_info=case["data_source_info"],
)
# Test data_source_info_dict (raw data)
raw_result = document.data_source_info_dict
assert raw_result == case["expected_raw"], f"Failed for {case['data_source_type']}"
# Verify raw_result is always a valid dict
assert isinstance(raw_result, dict)
def test_local_file_data_source_info_without_db_context(self):
"""Test that local_file type data_source_info_dict works without database context."""
test_data: LocalFileInfo = {
"file_path": "/local/path/document.txt",
"size": 512,
"created_at": "2024-01-01T00:00:00Z",
}
document = Document(
data_source_type="local_file",
data_source_info=json.dumps(test_data),
)
# data_source_info_dict should return the raw data (this doesn't need DB context)
raw_data = document.data_source_info_dict
assert raw_data == test_data
assert isinstance(raw_data, dict)
# Verify the data contains expected keys for pipeline mode
assert "file_path" in raw_data
assert "size" in raw_data
def test_notion_and_website_crawl_data_source_detail(self):
"""Test that notion_import and website_crawl return raw data in data_source_detail_dict."""
# Test notion_import
notion_data: NotionImportInfo = {"notion_page_id": "page-123", "workspace_id": "ws-456"}
document = Document(
data_source_type="notion_import",
data_source_info=json.dumps(notion_data),
)
# data_source_detail_dict should return raw data for notion_import
detail_result = document.data_source_detail_dict
assert detail_result == notion_data
# Test website_crawl
website_data: WebsiteCrawlInfo = {"url": "https://example.com", "job_id": "job-789"}
document = Document(
data_source_type="website_crawl",
data_source_info=json.dumps(website_data),
)
# data_source_detail_dict should return raw data for website_crawl
detail_result = document.data_source_detail_dict
assert detail_result == website_data
def test_local_file_data_source_detail_dict_without_db(self):
"""Test that local_file returns empty data_source_detail_dict (this doesn't need DB context)."""
# Test local_file - this should work without database context since it returns {} early
document = Document(
data_source_type="local_file",
data_source_info=json.dumps({"file_path": "/tmp/test.txt"}),
)
# Should return empty dict for local_file type (handled in the model)
detail_result = document.data_source_detail_dict
assert detail_result == {}

View File

@ -0,0 +1,27 @@
import builtins
import pytest
from flask import Flask
from flask.views import MethodView
from extensions import ext_fastopenapi
if not hasattr(builtins, "MethodView"):
builtins.MethodView = MethodView # type: ignore[attr-defined]
@pytest.fixture
def app() -> Flask:
app = Flask(__name__)
app.config["TESTING"] = True
return app
def test_console_ping_fastopenapi_returns_pong(app: Flask):
ext_fastopenapi.init_app(app)
client = app.test_client()
response = client.get("/console/api/ping")
assert response.status_code == 200
assert response.get_json() == {"result": "pong"}

View File

@ -0,0 +1,56 @@
import builtins
from unittest.mock import patch
import pytest
from flask import Flask
from flask.views import MethodView
from extensions import ext_fastopenapi
if not hasattr(builtins, "MethodView"):
builtins.MethodView = MethodView # type: ignore[attr-defined]
@pytest.fixture
def app() -> Flask:
app = Flask(__name__)
app.config["TESTING"] = True
return app
def test_console_setup_fastopenapi_get_not_started(app: Flask):
ext_fastopenapi.init_app(app)
with (
patch("controllers.console.setup.dify_config.EDITION", "SELF_HOSTED"),
patch("controllers.console.setup.get_setup_status", return_value=None),
):
client = app.test_client()
response = client.get("/console/api/setup")
assert response.status_code == 200
assert response.get_json() == {"step": "not_started", "setup_at": None}
def test_console_setup_fastopenapi_post_success(app: Flask):
ext_fastopenapi.init_app(app)
payload = {
"email": "admin@example.com",
"name": "Admin",
"password": "Passw0rd1",
"language": "en-US",
}
with (
patch("controllers.console.wraps.dify_config.EDITION", "SELF_HOSTED"),
patch("controllers.console.setup.get_setup_status", return_value=None),
patch("controllers.console.setup.TenantService.get_tenant_count", return_value=0),
patch("controllers.console.setup.get_init_validate_status", return_value=True),
patch("controllers.console.setup.RegisterService.setup"),
):
client = app.test_client()
response = client.post("/console/api/setup", json=payload)
assert response.status_code == 201
assert response.get_json() == {"result": "success"}

View File

@ -0,0 +1,35 @@
import builtins
from unittest.mock import patch
import pytest
from flask import Flask
from flask.views import MethodView
from configs import dify_config
from extensions import ext_fastopenapi
if not hasattr(builtins, "MethodView"):
builtins.MethodView = MethodView # type: ignore[attr-defined]
@pytest.fixture
def app() -> Flask:
app = Flask(__name__)
app.config["TESTING"] = True
return app
def test_console_version_fastopenapi_returns_current_version(app: Flask):
ext_fastopenapi.init_app(app)
with patch("controllers.console.version.dify_config.CHECK_UPDATE_URL", None):
client = app.test_client()
response = client.get("/console/api/version", query_string={"current_version": "0.0.0"})
assert response.status_code == 200
data = response.get_json()
assert data["version"] == dify_config.project.version
assert data["release_date"] == ""
assert data["release_notes"] == ""
assert data["can_auto_update"] is False
assert "features" in data

View File

@ -1,7 +1,9 @@
import builtins
import io
from unittest.mock import patch
import pytest
from flask.views import MethodView
from werkzeug.exceptions import Forbidden
from controllers.common.errors import (
@ -14,6 +16,9 @@ from controllers.common.errors import (
from services.errors.file import FileTooLargeError as ServiceFileTooLargeError
from services.errors.file import UnsupportedFileTypeError as ServiceUnsupportedFileTypeError
if not hasattr(builtins, "MethodView"):
builtins.MethodView = MethodView # type: ignore[attr-defined]
class TestFileUploadSecurity:
"""Test file upload security logic without complex framework setup"""
@ -128,7 +133,7 @@ class TestFileUploadSecurity:
# Test passes if no exception is raised
# Test 4: Service error handling
@patch("services.file_service.FileService.upload_file")
@patch("controllers.console.files.FileService.upload_file")
def test_should_handle_file_too_large_error(self, mock_upload):
"""Test that service FileTooLargeError is properly converted"""
mock_upload.side_effect = ServiceFileTooLargeError("File too large")
@ -140,7 +145,7 @@ class TestFileUploadSecurity:
with pytest.raises(FileTooLargeError):
raise FileTooLargeError(e.description)
@patch("services.file_service.FileService.upload_file")
@patch("controllers.console.files.FileService.upload_file")
def test_should_handle_unsupported_file_type_error(self, mock_upload):
"""Test that service UnsupportedFileTypeError is properly converted"""
mock_upload.side_effect = ServiceUnsupportedFileTypeError()

View File

@ -0,0 +1,247 @@
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from flask import Flask, g
from controllers.console.workspace.account import (
AccountDeleteUpdateFeedbackApi,
ChangeEmailCheckApi,
ChangeEmailResetApi,
ChangeEmailSendEmailApi,
CheckEmailUnique,
)
from models import Account
from services.account_service import AccountService
@pytest.fixture
def app():
app = Flask(__name__)
app.config["TESTING"] = True
app.config["RESTX_MASK_HEADER"] = "X-Fields"
app.login_manager = SimpleNamespace(_load_user=lambda: None)
return app
def _mock_wraps_db(mock_db):
mock_db.session.query.return_value.first.return_value = MagicMock()
def _build_account(email: str, account_id: str = "acc", tenant: object | None = None) -> Account:
tenant_obj = tenant if tenant is not None else SimpleNamespace(id="tenant-id")
account = Account(name=account_id, email=email)
account.email = email
account.id = account_id
account.status = "active"
account._current_tenant = tenant_obj
return account
def _set_logged_in_user(account: Account):
g._login_user = account
g._current_tenant = account.current_tenant
class TestChangeEmailSend:
@patch("controllers.console.wraps.db")
@patch("controllers.console.workspace.account.current_account_with_tenant")
@patch("controllers.console.workspace.account.AccountService.get_change_email_data")
@patch("controllers.console.workspace.account.AccountService.send_change_email_email")
@patch("controllers.console.workspace.account.AccountService.is_email_send_ip_limit", return_value=False)
@patch("controllers.console.workspace.account.extract_remote_ip", return_value="127.0.0.1")
@patch("libs.login.check_csrf_token", return_value=None)
@patch("controllers.console.wraps.FeatureService.get_system_features")
def test_should_normalize_new_email_phase(
self,
mock_features,
mock_csrf,
mock_extract_ip,
mock_is_ip_limit,
mock_send_email,
mock_get_change_data,
mock_current_account,
mock_db,
app,
):
_mock_wraps_db(mock_db)
mock_features.return_value = SimpleNamespace(enable_change_email=True)
mock_account = _build_account("current@example.com", "acc1")
mock_current_account.return_value = (mock_account, None)
mock_get_change_data.return_value = {"email": "current@example.com"}
mock_send_email.return_value = "token-abc"
with app.test_request_context(
"/account/change-email",
method="POST",
json={"email": "New@Example.com", "language": "en-US", "phase": "new_email", "token": "token-123"},
):
_set_logged_in_user(_build_account("tester@example.com", "tester"))
response = ChangeEmailSendEmailApi().post()
assert response == {"result": "success", "data": "token-abc"}
mock_send_email.assert_called_once_with(
account=None,
email="new@example.com",
old_email="current@example.com",
language="en-US",
phase="new_email",
)
mock_extract_ip.assert_called_once()
mock_is_ip_limit.assert_called_once_with("127.0.0.1")
mock_csrf.assert_called_once()
class TestChangeEmailValidity:
@patch("controllers.console.wraps.db")
@patch("controllers.console.workspace.account.current_account_with_tenant")
@patch("controllers.console.workspace.account.AccountService.reset_change_email_error_rate_limit")
@patch("controllers.console.workspace.account.AccountService.generate_change_email_token")
@patch("controllers.console.workspace.account.AccountService.revoke_change_email_token")
@patch("controllers.console.workspace.account.AccountService.add_change_email_error_rate_limit")
@patch("controllers.console.workspace.account.AccountService.get_change_email_data")
@patch("controllers.console.workspace.account.AccountService.is_change_email_error_rate_limit")
@patch("libs.login.check_csrf_token", return_value=None)
@patch("controllers.console.wraps.FeatureService.get_system_features")
def test_should_validate_with_normalized_email(
self,
mock_features,
mock_csrf,
mock_is_rate_limit,
mock_get_data,
mock_add_rate,
mock_revoke_token,
mock_generate_token,
mock_reset_rate,
mock_current_account,
mock_db,
app,
):
_mock_wraps_db(mock_db)
mock_features.return_value = SimpleNamespace(enable_change_email=True)
mock_account = _build_account("user@example.com", "acc2")
mock_current_account.return_value = (mock_account, None)
mock_is_rate_limit.return_value = False
mock_get_data.return_value = {"email": "user@example.com", "code": "1234", "old_email": "old@example.com"}
mock_generate_token.return_value = (None, "new-token")
with app.test_request_context(
"/account/change-email/validity",
method="POST",
json={"email": "User@Example.com", "code": "1234", "token": "token-123"},
):
_set_logged_in_user(_build_account("tester@example.com", "tester"))
response = ChangeEmailCheckApi().post()
assert response == {"is_valid": True, "email": "user@example.com", "token": "new-token"}
mock_is_rate_limit.assert_called_once_with("user@example.com")
mock_add_rate.assert_not_called()
mock_revoke_token.assert_called_once_with("token-123")
mock_generate_token.assert_called_once_with(
"user@example.com", code="1234", old_email="old@example.com", additional_data={}
)
mock_reset_rate.assert_called_once_with("user@example.com")
mock_csrf.assert_called_once()
class TestChangeEmailReset:
@patch("controllers.console.wraps.db")
@patch("controllers.console.workspace.account.current_account_with_tenant")
@patch("controllers.console.workspace.account.AccountService.send_change_email_completed_notify_email")
@patch("controllers.console.workspace.account.AccountService.update_account_email")
@patch("controllers.console.workspace.account.AccountService.revoke_change_email_token")
@patch("controllers.console.workspace.account.AccountService.get_change_email_data")
@patch("controllers.console.workspace.account.AccountService.check_email_unique")
@patch("controllers.console.workspace.account.AccountService.is_account_in_freeze")
@patch("libs.login.check_csrf_token", return_value=None)
@patch("controllers.console.wraps.FeatureService.get_system_features")
def test_should_normalize_new_email_before_update(
self,
mock_features,
mock_csrf,
mock_is_freeze,
mock_check_unique,
mock_get_data,
mock_revoke_token,
mock_update_account,
mock_send_notify,
mock_current_account,
mock_db,
app,
):
_mock_wraps_db(mock_db)
mock_features.return_value = SimpleNamespace(enable_change_email=True)
current_user = _build_account("old@example.com", "acc3")
mock_current_account.return_value = (current_user, None)
mock_is_freeze.return_value = False
mock_check_unique.return_value = True
mock_get_data.return_value = {"old_email": "OLD@example.com"}
mock_account_after_update = _build_account("new@example.com", "acc3-updated")
mock_update_account.return_value = mock_account_after_update
with app.test_request_context(
"/account/change-email/reset",
method="POST",
json={"new_email": "New@Example.com", "token": "token-123"},
):
_set_logged_in_user(_build_account("tester@example.com", "tester"))
ChangeEmailResetApi().post()
mock_is_freeze.assert_called_once_with("new@example.com")
mock_check_unique.assert_called_once_with("new@example.com")
mock_revoke_token.assert_called_once_with("token-123")
mock_update_account.assert_called_once_with(current_user, email="new@example.com")
mock_send_notify.assert_called_once_with(email="new@example.com")
mock_csrf.assert_called_once()
class TestAccountDeletionFeedback:
@patch("controllers.console.wraps.db")
@patch("controllers.console.workspace.account.BillingService.update_account_deletion_feedback")
def test_should_normalize_feedback_email(self, mock_update, mock_db, app):
_mock_wraps_db(mock_db)
with app.test_request_context(
"/account/delete/feedback",
method="POST",
json={"email": "User@Example.com", "feedback": "test"},
):
response = AccountDeleteUpdateFeedbackApi().post()
assert response == {"result": "success"}
mock_update.assert_called_once_with("User@Example.com", "test")
class TestCheckEmailUnique:
@patch("controllers.console.wraps.db")
@patch("controllers.console.workspace.account.AccountService.check_email_unique")
@patch("controllers.console.workspace.account.AccountService.is_account_in_freeze")
def test_should_normalize_email(self, mock_is_freeze, mock_check_unique, mock_db, app):
_mock_wraps_db(mock_db)
mock_is_freeze.return_value = False
mock_check_unique.return_value = True
with app.test_request_context(
"/account/change-email/check-email-unique",
method="POST",
json={"email": "Case@Test.com"},
):
response = CheckEmailUnique().post()
assert response == {"result": "success"}
mock_is_freeze.assert_called_once_with("case@test.com")
mock_check_unique.assert_called_once_with("case@test.com")
def test_get_account_by_email_with_case_fallback_uses_lowercase_lookup():
session = MagicMock()
first = MagicMock()
first.scalar_one_or_none.return_value = None
second = MagicMock()
expected_account = MagicMock()
second.scalar_one_or_none.return_value = expected_account
session.execute.side_effect = [first, second]
result = AccountService.get_account_by_email_with_case_fallback("Mixed@Test.com", session=session)
assert result is expected_account
assert session.execute.call_count == 2

View File

@ -0,0 +1,82 @@
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from flask import Flask, g
from controllers.console.workspace.members import MemberInviteEmailApi
from models.account import Account, TenantAccountRole
@pytest.fixture
def app():
flask_app = Flask(__name__)
flask_app.config["TESTING"] = True
flask_app.login_manager = SimpleNamespace(_load_user=lambda: None)
return flask_app
def _mock_wraps_db(mock_db):
mock_db.session.query.return_value.first.return_value = MagicMock()
def _build_feature_flags():
placeholder_quota = SimpleNamespace(limit=0, size=0)
workspace_members = SimpleNamespace(is_available=lambda count: True)
return SimpleNamespace(
billing=SimpleNamespace(enabled=False),
workspace_members=workspace_members,
members=placeholder_quota,
apps=placeholder_quota,
vector_space=placeholder_quota,
documents_upload_quota=placeholder_quota,
annotation_quota_limit=placeholder_quota,
)
class TestMemberInviteEmailApi:
@patch("controllers.console.workspace.members.FeatureService.get_features")
@patch("controllers.console.workspace.members.RegisterService.invite_new_member")
@patch("controllers.console.workspace.members.current_account_with_tenant")
@patch("controllers.console.wraps.db")
@patch("libs.login.check_csrf_token", return_value=None)
def test_invite_normalizes_emails(
self,
mock_csrf,
mock_db,
mock_current_account,
mock_invite_member,
mock_get_features,
app,
):
_mock_wraps_db(mock_db)
mock_get_features.return_value = _build_feature_flags()
mock_invite_member.return_value = "token-abc"
tenant = SimpleNamespace(id="tenant-1", name="Test Tenant")
inviter = SimpleNamespace(email="Owner@Example.com", current_tenant=tenant, status="active")
mock_current_account.return_value = (inviter, tenant.id)
with patch("controllers.console.workspace.members.dify_config.CONSOLE_WEB_URL", "https://console.example.com"):
with app.test_request_context(
"/workspaces/current/members/invite-email",
method="POST",
json={"emails": ["User@Example.com"], "role": TenantAccountRole.EDITOR.value, "language": "en-US"},
):
account = Account(name="tester", email="tester@example.com")
account._current_tenant = tenant
g._login_user = account
g._current_tenant = tenant
response, status_code = MemberInviteEmailApi().post()
assert status_code == 201
assert response["invitation_results"][0]["email"] == "user@example.com"
assert mock_invite_member.call_count == 1
call_args = mock_invite_member.call_args
assert call_args.kwargs["tenant"] == tenant
assert call_args.kwargs["email"] == "User@Example.com"
assert call_args.kwargs["language"] == "en-US"
assert call_args.kwargs["role"] == TenantAccountRole.EDITOR
assert call_args.kwargs["inviter"] == inviter
mock_csrf.assert_called_once()

View File

@ -0,0 +1,145 @@
"""Unit tests for load balancing credential validation APIs."""
from __future__ import annotations
import builtins
import importlib
import sys
from types import SimpleNamespace
from unittest.mock import MagicMock
import pytest
from flask import Flask
from flask.views import MethodView
from werkzeug.exceptions import Forbidden
from core.model_runtime.entities.model_entities import ModelType
from core.model_runtime.errors.validate import CredentialsValidateFailedError
if not hasattr(builtins, "MethodView"):
builtins.MethodView = MethodView # type: ignore[attr-defined]
from models.account import TenantAccountRole
@pytest.fixture
def app() -> Flask:
app = Flask(__name__)
app.config["TESTING"] = True
return app
@pytest.fixture
def load_balancing_module(monkeypatch: pytest.MonkeyPatch):
"""Reload controller module with lightweight decorators for testing."""
from controllers.console import console_ns, wraps
from libs import login
def _noop(func):
return func
monkeypatch.setattr(login, "login_required", _noop)
monkeypatch.setattr(wraps, "setup_required", _noop)
monkeypatch.setattr(wraps, "account_initialization_required", _noop)
def _noop_route(*args, **kwargs): # type: ignore[override]
def _decorator(cls):
return cls
return _decorator
monkeypatch.setattr(console_ns, "route", _noop_route)
module_name = "controllers.console.workspace.load_balancing_config"
sys.modules.pop(module_name, None)
module = importlib.import_module(module_name)
return module
def _mock_user(role: TenantAccountRole) -> SimpleNamespace:
return SimpleNamespace(current_role=role)
def _prepare_context(module, monkeypatch: pytest.MonkeyPatch, role=TenantAccountRole.OWNER):
user = _mock_user(role)
monkeypatch.setattr(module, "current_account_with_tenant", lambda: (user, "tenant-123"))
mock_service = MagicMock()
monkeypatch.setattr(module, "ModelLoadBalancingService", lambda: mock_service)
return mock_service
def _request_payload():
return {"model": "gpt-4o", "model_type": ModelType.LLM, "credentials": {"api_key": "sk-***"}}
def test_validate_credentials_success(app: Flask, load_balancing_module, monkeypatch: pytest.MonkeyPatch):
service = _prepare_context(load_balancing_module, monkeypatch)
with app.test_request_context(
"/workspaces/current/model-providers/openai/models/load-balancing-configs/credentials-validate",
method="POST",
json=_request_payload(),
):
response = load_balancing_module.LoadBalancingCredentialsValidateApi().post(provider="openai")
assert response == {"result": "success"}
service.validate_load_balancing_credentials.assert_called_once_with(
tenant_id="tenant-123",
provider="openai",
model="gpt-4o",
model_type=ModelType.LLM,
credentials={"api_key": "sk-***"},
)
def test_validate_credentials_returns_error_message(app: Flask, load_balancing_module, monkeypatch: pytest.MonkeyPatch):
service = _prepare_context(load_balancing_module, monkeypatch)
service.validate_load_balancing_credentials.side_effect = CredentialsValidateFailedError("invalid credentials")
with app.test_request_context(
"/workspaces/current/model-providers/openai/models/load-balancing-configs/credentials-validate",
method="POST",
json=_request_payload(),
):
response = load_balancing_module.LoadBalancingCredentialsValidateApi().post(provider="openai")
assert response == {"result": "error", "error": "invalid credentials"}
def test_validate_credentials_requires_privileged_role(
app: Flask, load_balancing_module, monkeypatch: pytest.MonkeyPatch
):
_prepare_context(load_balancing_module, monkeypatch, role=TenantAccountRole.NORMAL)
with app.test_request_context(
"/workspaces/current/model-providers/openai/models/load-balancing-configs/credentials-validate",
method="POST",
json=_request_payload(),
):
api = load_balancing_module.LoadBalancingCredentialsValidateApi()
with pytest.raises(Forbidden):
api.post(provider="openai")
def test_validate_credentials_with_config_id(app: Flask, load_balancing_module, monkeypatch: pytest.MonkeyPatch):
service = _prepare_context(load_balancing_module, monkeypatch)
with app.test_request_context(
"/workspaces/current/model-providers/openai/models/load-balancing-configs/cfg-1/credentials-validate",
method="POST",
json=_request_payload(),
):
response = load_balancing_module.LoadBalancingConfigCredentialsValidateApi().post(
provider="openai", config_id="cfg-1"
)
assert response == {"result": "success"}
service.validate_load_balancing_credentials.assert_called_once_with(
tenant_id="tenant-123",
provider="openai",
model="gpt-4o",
model_type=ModelType.LLM,
credentials={"api_key": "sk-***"},
config_id="cfg-1",
)

View File

@ -0,0 +1,100 @@
import json
from unittest.mock import MagicMock, patch
import pytest
from flask import Flask
from flask_restx import Api
from controllers.console.workspace.tool_providers import ToolProviderMCPApi
from core.db.session_factory import configure_session_factory
from extensions.ext_database import db
from services.tools.mcp_tools_manage_service import ReconnectResult
# Backward-compat fixtures referenced by @pytest.mark.usefixtures in this file.
# They are intentionally no-ops because the test already patches the required
# behaviors explicitly via @patch and context managers below.
@pytest.fixture
def _mock_cache():
return
@pytest.fixture
def _mock_user_tenant():
return
@pytest.fixture
def client():
app = Flask(__name__)
app.config["TESTING"] = True
app.config["SQLALCHEMY_DATABASE_URI"] = "sqlite:///:memory:"
api = Api(app)
api.add_resource(ToolProviderMCPApi, "/console/api/workspaces/current/tool-provider/mcp")
db.init_app(app)
# Configure session factory used by controller code
with app.app_context():
configure_session_factory(db.engine)
return app.test_client()
@patch(
"controllers.console.workspace.tool_providers.current_account_with_tenant", return_value=(MagicMock(id="u1"), "t1")
)
@patch("controllers.console.workspace.tool_providers.Session")
@patch("controllers.console.workspace.tool_providers.MCPToolManageService._reconnect_with_url")
@pytest.mark.usefixtures("_mock_cache", "_mock_user_tenant")
def test_create_mcp_provider_populates_tools(mock_reconnect, mock_session, mock_current_account_with_tenant, client):
# Arrange: reconnect returns tools immediately
mock_reconnect.return_value = ReconnectResult(
authed=True,
tools=json.dumps(
[{"name": "ping", "description": "ok", "inputSchema": {"type": "object"}, "outputSchema": {}}]
),
encrypted_credentials="{}",
)
# Fake service.create_provider -> returns object with id for reload
svc = MagicMock()
create_result = MagicMock()
create_result.id = "provider-1"
svc.create_provider.return_value = create_result
svc.get_provider.return_value = MagicMock(id="provider-1", tenant_id="t1") # used by reload path
mock_session.return_value.__enter__.return_value = MagicMock()
# Patch MCPToolManageService constructed inside controller
with patch("controllers.console.workspace.tool_providers.MCPToolManageService", return_value=svc):
payload = {
"server_url": "http://example.com/mcp",
"name": "demo",
"icon": "😀",
"icon_type": "emoji",
"icon_background": "#000",
"server_identifier": "demo-sid",
"configuration": {"timeout": 5, "sse_read_timeout": 30},
"headers": {},
"authentication": {},
}
# Act
with (
patch("controllers.console.wraps.dify_config.EDITION", "CLOUD"), # bypass setup_required DB check
patch("controllers.console.wraps.current_account_with_tenant", return_value=(MagicMock(id="u1"), "t1")),
patch("libs.login.check_csrf_token", return_value=None), # bypass CSRF in login_required
patch("libs.login._get_user", return_value=MagicMock(id="u1", is_authenticated=True)), # login
patch(
"services.tools.tools_transform_service.ToolTransformService.mcp_provider_to_user_provider",
return_value={"id": "provider-1", "tools": [{"name": "ping"}]},
),
):
resp = client.post(
"/console/api/workspaces/current/tool-provider/mcp",
data=json.dumps(payload),
content_type="application/json",
)
# Assert
assert resp.status_code == 200
body = resp.get_json()
assert body.get("id") == "provider-1"
# 若 transform 后包含 tools 字段,确保非空
assert isinstance(body.get("tools"), list)
assert body["tools"]

View File

@ -1,195 +0,0 @@
"""Unit tests for controllers.web.forgot_password endpoints."""
from __future__ import annotations
import base64
import builtins
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from flask import Flask
from flask.views import MethodView
# Ensure flask_restx.api finds MethodView during import.
if not hasattr(builtins, "MethodView"):
builtins.MethodView = MethodView # type: ignore[attr-defined]
def _load_controller_module():
"""Import controllers.web.forgot_password using a stub package."""
import importlib
import importlib.util
import sys
from types import ModuleType
parent_module_name = "controllers.web"
module_name = f"{parent_module_name}.forgot_password"
if parent_module_name not in sys.modules:
from flask_restx import Namespace
stub = ModuleType(parent_module_name)
stub.__file__ = "controllers/web/__init__.py"
stub.__path__ = ["controllers/web"]
stub.__package__ = "controllers"
stub.__spec__ = importlib.util.spec_from_loader(parent_module_name, loader=None, is_package=True)
stub.web_ns = Namespace("web", description="Web API", path="/")
sys.modules[parent_module_name] = stub
return importlib.import_module(module_name)
forgot_password_module = _load_controller_module()
ForgotPasswordCheckApi = forgot_password_module.ForgotPasswordCheckApi
ForgotPasswordResetApi = forgot_password_module.ForgotPasswordResetApi
ForgotPasswordSendEmailApi = forgot_password_module.ForgotPasswordSendEmailApi
@pytest.fixture
def app() -> Flask:
"""Configure a minimal Flask app for request contexts."""
app = Flask(__name__)
app.config["TESTING"] = True
return app
@pytest.fixture(autouse=True)
def _enable_web_endpoint_guards():
"""Stub enterprise and feature toggles used by route decorators."""
features = SimpleNamespace(enable_email_password_login=True)
with (
patch("controllers.console.wraps.dify_config.ENTERPRISE_ENABLED", True),
patch("controllers.console.wraps.dify_config.EDITION", "CLOUD"),
patch("controllers.console.wraps.FeatureService.get_system_features", return_value=features),
):
yield
@pytest.fixture(autouse=True)
def _mock_controller_db():
"""Replace controller-level db reference with a simple stub."""
fake_db = SimpleNamespace(engine=MagicMock(name="engine"))
fake_wraps_db = SimpleNamespace(
session=MagicMock(query=MagicMock(return_value=MagicMock(first=MagicMock(return_value=True))))
)
with (
patch("controllers.web.forgot_password.db", fake_db),
patch("controllers.console.wraps.db", fake_wraps_db),
):
yield fake_db
@patch("controllers.web.forgot_password.AccountService.send_reset_password_email", return_value="reset-token")
@patch("controllers.web.forgot_password.Session")
@patch("controllers.web.forgot_password.AccountService.is_email_send_ip_limit", return_value=False)
@patch("controllers.web.forgot_password.extract_remote_ip", return_value="203.0.113.10")
def test_send_reset_email_success(
mock_extract_ip: MagicMock,
mock_is_ip_limit: MagicMock,
mock_session: MagicMock,
mock_send_email: MagicMock,
app: Flask,
):
"""POST /forgot-password returns token when email exists and limits allow."""
mock_account = MagicMock()
session_ctx = MagicMock()
mock_session.return_value.__enter__.return_value = session_ctx
session_ctx.execute.return_value.scalar_one_or_none.return_value = mock_account
with app.test_request_context(
"/forgot-password",
method="POST",
json={"email": "user@example.com"},
):
response = ForgotPasswordSendEmailApi().post()
assert response == {"result": "success", "data": "reset-token"}
mock_extract_ip.assert_called_once()
mock_is_ip_limit.assert_called_once_with("203.0.113.10")
mock_send_email.assert_called_once_with(account=mock_account, email="user@example.com", language="en-US")
@patch("controllers.web.forgot_password.AccountService.reset_forgot_password_error_rate_limit")
@patch("controllers.web.forgot_password.AccountService.generate_reset_password_token", return_value=({}, "new-token"))
@patch("controllers.web.forgot_password.AccountService.revoke_reset_password_token")
@patch("controllers.web.forgot_password.AccountService.get_reset_password_data")
@patch("controllers.web.forgot_password.AccountService.is_forgot_password_error_rate_limit", return_value=False)
def test_check_token_success(
mock_is_rate_limited: MagicMock,
mock_get_data: MagicMock,
mock_revoke: MagicMock,
mock_generate: MagicMock,
mock_reset_limit: MagicMock,
app: Flask,
):
"""POST /forgot-password/validity validates the code and refreshes token."""
mock_get_data.return_value = {"email": "user@example.com", "code": "123456"}
with app.test_request_context(
"/forgot-password/validity",
method="POST",
json={"email": "user@example.com", "code": "123456", "token": "old-token"},
):
response = ForgotPasswordCheckApi().post()
assert response == {"is_valid": True, "email": "user@example.com", "token": "new-token"}
mock_is_rate_limited.assert_called_once_with("user@example.com")
mock_get_data.assert_called_once_with("old-token")
mock_revoke.assert_called_once_with("old-token")
mock_generate.assert_called_once_with(
"user@example.com",
code="123456",
additional_data={"phase": "reset"},
)
mock_reset_limit.assert_called_once_with("user@example.com")
@patch("controllers.web.forgot_password.hash_password", return_value=b"hashed-value")
@patch("controllers.web.forgot_password.secrets.token_bytes", return_value=b"0123456789abcdef")
@patch("controllers.web.forgot_password.Session")
@patch("controllers.web.forgot_password.AccountService.revoke_reset_password_token")
@patch("controllers.web.forgot_password.AccountService.get_reset_password_data")
def test_reset_password_success(
mock_get_data: MagicMock,
mock_revoke_token: MagicMock,
mock_session: MagicMock,
mock_token_bytes: MagicMock,
mock_hash_password: MagicMock,
app: Flask,
):
"""POST /forgot-password/resets updates the stored password when token is valid."""
mock_get_data.return_value = {"email": "user@example.com", "phase": "reset"}
account = MagicMock()
session_ctx = MagicMock()
mock_session.return_value.__enter__.return_value = session_ctx
session_ctx.execute.return_value.scalar_one_or_none.return_value = account
with app.test_request_context(
"/forgot-password/resets",
method="POST",
json={
"token": "reset-token",
"new_password": "StrongPass123!",
"password_confirm": "StrongPass123!",
},
):
response = ForgotPasswordResetApi().post()
assert response == {"result": "success"}
mock_get_data.assert_called_once_with("reset-token")
mock_revoke_token.assert_called_once_with("reset-token")
mock_token_bytes.assert_called_once_with(16)
mock_hash_password.assert_called_once_with("StrongPass123!", b"0123456789abcdef")
expected_password = base64.b64encode(b"hashed-value").decode()
assert account.password == expected_password
expected_salt = base64.b64encode(b"0123456789abcdef").decode()
assert account.password_salt == expected_salt
session_ctx.commit.assert_called_once()

View File

@ -0,0 +1,174 @@
"""Unit tests for controllers.web.message message list mapping."""
from __future__ import annotations
import builtins
from datetime import datetime
from types import ModuleType, SimpleNamespace
from unittest.mock import patch
from uuid import uuid4
import pytest
from flask import Flask
from flask.views import MethodView
# Ensure flask_restx.api finds MethodView during import.
if not hasattr(builtins, "MethodView"):
builtins.MethodView = MethodView # type: ignore[attr-defined]
def _load_controller_module():
"""Import controllers.web.message using a stub package."""
import importlib
import importlib.util
import sys
parent_module_name = "controllers.web"
module_name = f"{parent_module_name}.message"
if parent_module_name not in sys.modules:
from flask_restx import Namespace
stub = ModuleType(parent_module_name)
stub.__file__ = "controllers/web/__init__.py"
stub.__path__ = ["controllers/web"]
stub.__package__ = "controllers"
stub.__spec__ = importlib.util.spec_from_loader(parent_module_name, loader=None, is_package=True)
stub.web_ns = Namespace("web", description="Web API", path="/")
sys.modules[parent_module_name] = stub
wraps_module_name = f"{parent_module_name}.wraps"
if wraps_module_name not in sys.modules:
wraps_stub = ModuleType(wraps_module_name)
class WebApiResource:
pass
wraps_stub.WebApiResource = WebApiResource
sys.modules[wraps_module_name] = wraps_stub
return importlib.import_module(module_name)
message_module = _load_controller_module()
MessageListApi = message_module.MessageListApi
@pytest.fixture
def app() -> Flask:
app = Flask(__name__)
app.config["TESTING"] = True
return app
def test_message_list_mapping(app: Flask) -> None:
conversation_id = str(uuid4())
message_id = str(uuid4())
created_at = datetime(2024, 1, 1, 12, 0, 0)
resource_created_at = datetime(2024, 1, 1, 13, 0, 0)
thought_created_at = datetime(2024, 1, 1, 14, 0, 0)
retriever_resource_obj = SimpleNamespace(
id="res-obj",
message_id=message_id,
position=2,
dataset_id="ds-1",
dataset_name="dataset",
document_id="doc-1",
document_name="document",
data_source_type="file",
segment_id="seg-1",
score=0.9,
hit_count=1,
word_count=10,
segment_position=0,
index_node_hash="hash",
content="content",
created_at=resource_created_at,
)
agent_thought = SimpleNamespace(
id="thought-1",
chain_id=None,
message_chain_id="chain-1",
message_id=message_id,
position=1,
thought="thinking",
tool="tool",
tool_labels={"label": "value"},
tool_input="{}",
created_at=thought_created_at,
observation="observed",
files=["file-a"],
)
message_file_obj = SimpleNamespace(
id="file-obj",
filename="b.txt",
type="file",
url=None,
mime_type=None,
size=None,
transfer_method="local",
belongs_to=None,
upload_file_id=None,
)
message = SimpleNamespace(
id=message_id,
conversation_id=conversation_id,
parent_message_id=None,
inputs={"foo": "bar"},
query="hello",
re_sign_file_url_answer="answer",
user_feedback=SimpleNamespace(rating="like"),
retriever_resources=[
{"id": "res-dict", "message_id": message_id, "position": 1},
retriever_resource_obj,
],
created_at=created_at,
agent_thoughts=[agent_thought],
message_files=[
{"id": "file-dict", "filename": "a.txt", "type": "file", "transfer_method": "local"},
message_file_obj,
],
status="success",
error=None,
message_metadata_dict={"meta": "value"},
)
pagination = SimpleNamespace(limit=20, has_more=False, data=[message])
app_model = SimpleNamespace(mode="chat")
end_user = SimpleNamespace()
with (
patch.object(message_module.MessageService, "pagination_by_first_id", return_value=pagination) as mock_page,
app.test_request_context(f"/messages?conversation_id={conversation_id}&limit=20"),
):
response = MessageListApi().get(app_model, end_user)
mock_page.assert_called_once_with(app_model, end_user, conversation_id, None, 20)
assert response["limit"] == 20
assert response["has_more"] is False
assert len(response["data"]) == 1
item = response["data"][0]
assert item["id"] == message_id
assert item["conversation_id"] == conversation_id
assert item["inputs"] == {"foo": "bar"}
assert item["answer"] == "answer"
assert item["feedback"]["rating"] == "like"
assert item["metadata"] == {"meta": "value"}
assert item["created_at"] == int(created_at.timestamp())
assert item["retriever_resources"][0]["id"] == "res-dict"
assert item["retriever_resources"][1]["id"] == "res-obj"
assert item["retriever_resources"][1]["created_at"] == int(resource_created_at.timestamp())
assert item["agent_thoughts"][0]["chain_id"] == "chain-1"
assert item["agent_thoughts"][0]["created_at"] == int(thought_created_at.timestamp())
assert item["message_files"][0]["id"] == "file-dict"
assert item["message_files"][1]["id"] == "file-obj"

View File

@ -0,0 +1,226 @@
import base64
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from flask import Flask
from controllers.web.forgot_password import (
ForgotPasswordCheckApi,
ForgotPasswordResetApi,
ForgotPasswordSendEmailApi,
)
@pytest.fixture
def app():
flask_app = Flask(__name__)
flask_app.config["TESTING"] = True
return flask_app
@pytest.fixture(autouse=True)
def _patch_wraps():
wraps_features = SimpleNamespace(enable_email_password_login=True)
dify_settings = SimpleNamespace(ENTERPRISE_ENABLED=True, EDITION="CLOUD")
with (
patch("controllers.console.wraps.db") as mock_db,
patch("controllers.console.wraps.dify_config", dify_settings),
patch("controllers.console.wraps.FeatureService.get_system_features", return_value=wraps_features),
):
mock_db.session.query.return_value.first.return_value = MagicMock()
yield
class TestForgotPasswordSendEmailApi:
@patch("controllers.web.forgot_password.AccountService.send_reset_password_email")
@patch("controllers.web.forgot_password.AccountService.get_account_by_email_with_case_fallback")
@patch("controllers.web.forgot_password.AccountService.is_email_send_ip_limit", return_value=False)
@patch("controllers.web.forgot_password.extract_remote_ip", return_value="127.0.0.1")
@patch("controllers.web.forgot_password.Session")
def test_should_normalize_email_before_sending(
self,
mock_session_cls,
mock_extract_ip,
mock_rate_limit,
mock_get_account,
mock_send_mail,
app,
):
mock_account = MagicMock()
mock_get_account.return_value = mock_account
mock_send_mail.return_value = "token-123"
mock_session = MagicMock()
mock_session_cls.return_value.__enter__.return_value = mock_session
with patch("controllers.web.forgot_password.db", SimpleNamespace(engine="engine")):
with app.test_request_context(
"/web/forgot-password",
method="POST",
json={"email": "User@Example.com", "language": "zh-Hans"},
):
response = ForgotPasswordSendEmailApi().post()
assert response == {"result": "success", "data": "token-123"}
mock_get_account.assert_called_once_with("User@Example.com", session=mock_session)
mock_send_mail.assert_called_once_with(account=mock_account, email="user@example.com", language="zh-Hans")
mock_extract_ip.assert_called_once()
mock_rate_limit.assert_called_once_with("127.0.0.1")
class TestForgotPasswordCheckApi:
@patch("controllers.web.forgot_password.AccountService.reset_forgot_password_error_rate_limit")
@patch("controllers.web.forgot_password.AccountService.generate_reset_password_token")
@patch("controllers.web.forgot_password.AccountService.revoke_reset_password_token")
@patch("controllers.web.forgot_password.AccountService.add_forgot_password_error_rate_limit")
@patch("controllers.web.forgot_password.AccountService.get_reset_password_data")
@patch("controllers.web.forgot_password.AccountService.is_forgot_password_error_rate_limit")
def test_should_normalize_email_for_validity_checks(
self,
mock_is_rate_limit,
mock_get_data,
mock_add_rate,
mock_revoke_token,
mock_generate_token,
mock_reset_rate,
app,
):
mock_is_rate_limit.return_value = False
mock_get_data.return_value = {"email": "User@Example.com", "code": "1234"}
mock_generate_token.return_value = (None, "new-token")
with app.test_request_context(
"/web/forgot-password/validity",
method="POST",
json={"email": "User@Example.com", "code": "1234", "token": "token-123"},
):
response = ForgotPasswordCheckApi().post()
assert response == {"is_valid": True, "email": "user@example.com", "token": "new-token"}
mock_is_rate_limit.assert_called_once_with("user@example.com")
mock_add_rate.assert_not_called()
mock_revoke_token.assert_called_once_with("token-123")
mock_generate_token.assert_called_once_with(
"User@Example.com",
code="1234",
additional_data={"phase": "reset"},
)
mock_reset_rate.assert_called_once_with("user@example.com")
@patch("controllers.web.forgot_password.AccountService.reset_forgot_password_error_rate_limit")
@patch("controllers.web.forgot_password.AccountService.generate_reset_password_token")
@patch("controllers.web.forgot_password.AccountService.revoke_reset_password_token")
@patch("controllers.web.forgot_password.AccountService.get_reset_password_data")
@patch("controllers.web.forgot_password.AccountService.is_forgot_password_error_rate_limit")
def test_should_preserve_token_email_case(
self,
mock_is_rate_limit,
mock_get_data,
mock_revoke_token,
mock_generate_token,
mock_reset_rate,
app,
):
mock_is_rate_limit.return_value = False
mock_get_data.return_value = {"email": "MixedCase@Example.com", "code": "5678"}
mock_generate_token.return_value = (None, "fresh-token")
with app.test_request_context(
"/web/forgot-password/validity",
method="POST",
json={"email": "mixedcase@example.com", "code": "5678", "token": "token-upper"},
):
response = ForgotPasswordCheckApi().post()
assert response == {"is_valid": True, "email": "mixedcase@example.com", "token": "fresh-token"}
mock_generate_token.assert_called_once_with(
"MixedCase@Example.com",
code="5678",
additional_data={"phase": "reset"},
)
mock_revoke_token.assert_called_once_with("token-upper")
mock_reset_rate.assert_called_once_with("mixedcase@example.com")
class TestForgotPasswordResetApi:
@patch("controllers.web.forgot_password.ForgotPasswordResetApi._update_existing_account")
@patch("controllers.web.forgot_password.AccountService.get_account_by_email_with_case_fallback")
@patch("controllers.web.forgot_password.Session")
@patch("controllers.web.forgot_password.AccountService.revoke_reset_password_token")
@patch("controllers.web.forgot_password.AccountService.get_reset_password_data")
def test_should_fetch_account_with_fallback(
self,
mock_get_reset_data,
mock_revoke_token,
mock_session_cls,
mock_get_account,
mock_update_account,
app,
):
mock_get_reset_data.return_value = {"phase": "reset", "email": "User@Example.com", "code": "1234"}
mock_account = MagicMock()
mock_get_account.return_value = mock_account
mock_session = MagicMock()
mock_session_cls.return_value.__enter__.return_value = mock_session
with patch("controllers.web.forgot_password.db", SimpleNamespace(engine="engine")):
with app.test_request_context(
"/web/forgot-password/resets",
method="POST",
json={
"token": "token-123",
"new_password": "ValidPass123!",
"password_confirm": "ValidPass123!",
},
):
response = ForgotPasswordResetApi().post()
assert response == {"result": "success"}
mock_get_account.assert_called_once_with("User@Example.com", session=mock_session)
mock_update_account.assert_called_once()
mock_revoke_token.assert_called_once_with("token-123")
@patch("controllers.web.forgot_password.hash_password", return_value=b"hashed-value")
@patch("controllers.web.forgot_password.secrets.token_bytes", return_value=b"0123456789abcdef")
@patch("controllers.web.forgot_password.Session")
@patch("controllers.web.forgot_password.AccountService.revoke_reset_password_token")
@patch("controllers.web.forgot_password.AccountService.get_reset_password_data")
@patch("controllers.web.forgot_password.AccountService.get_account_by_email_with_case_fallback")
def test_should_update_password_and_commit(
self,
mock_get_account,
mock_get_reset_data,
mock_revoke_token,
mock_session_cls,
mock_token_bytes,
mock_hash_password,
app,
):
mock_get_reset_data.return_value = {"phase": "reset", "email": "user@example.com"}
account = MagicMock()
mock_get_account.return_value = account
mock_session = MagicMock()
mock_session_cls.return_value.__enter__.return_value = mock_session
with patch("controllers.web.forgot_password.db", SimpleNamespace(engine="engine")):
with app.test_request_context(
"/web/forgot-password/resets",
method="POST",
json={
"token": "reset-token",
"new_password": "StrongPass123!",
"password_confirm": "StrongPass123!",
},
):
response = ForgotPasswordResetApi().post()
assert response == {"result": "success"}
mock_get_reset_data.assert_called_once_with("reset-token")
mock_revoke_token.assert_called_once_with("reset-token")
mock_token_bytes.assert_called_once_with(16)
mock_hash_password.assert_called_once_with("StrongPass123!", b"0123456789abcdef")
expected_password = base64.b64encode(b"hashed-value").decode()
assert account.password == expected_password
expected_salt = base64.b64encode(b"0123456789abcdef").decode()
assert account.password_salt == expected_salt
mock_session.commit.assert_called_once()

View File

@ -0,0 +1,91 @@
import base64
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from flask import Flask
from controllers.web.login import EmailCodeLoginApi, EmailCodeLoginSendEmailApi
def encode_code(code: str) -> str:
return base64.b64encode(code.encode("utf-8")).decode()
@pytest.fixture
def app():
flask_app = Flask(__name__)
flask_app.config["TESTING"] = True
return flask_app
@pytest.fixture(autouse=True)
def _patch_wraps():
wraps_features = SimpleNamespace(enable_email_password_login=True)
console_dify = SimpleNamespace(ENTERPRISE_ENABLED=True, EDITION="CLOUD")
web_dify = SimpleNamespace(ENTERPRISE_ENABLED=True)
with (
patch("controllers.console.wraps.db") as mock_db,
patch("controllers.console.wraps.dify_config", console_dify),
patch("controllers.console.wraps.FeatureService.get_system_features", return_value=wraps_features),
patch("controllers.web.login.dify_config", web_dify),
):
mock_db.session.query.return_value.first.return_value = MagicMock()
yield
class TestEmailCodeLoginSendEmailApi:
@patch("controllers.web.login.WebAppAuthService.send_email_code_login_email")
@patch("controllers.web.login.WebAppAuthService.get_user_through_email")
def test_should_fetch_account_with_original_email(
self,
mock_get_user,
mock_send_email,
app,
):
mock_account = MagicMock()
mock_get_user.return_value = mock_account
mock_send_email.return_value = "token-123"
with app.test_request_context(
"/web/email-code-login",
method="POST",
json={"email": "User@Example.com", "language": "en-US"},
):
response = EmailCodeLoginSendEmailApi().post()
assert response == {"result": "success", "data": "token-123"}
mock_get_user.assert_called_once_with("User@Example.com")
mock_send_email.assert_called_once_with(account=mock_account, language="en-US")
class TestEmailCodeLoginApi:
@patch("controllers.web.login.AccountService.reset_login_error_rate_limit")
@patch("controllers.web.login.WebAppAuthService.login", return_value="new-access-token")
@patch("controllers.web.login.WebAppAuthService.get_user_through_email")
@patch("controllers.web.login.WebAppAuthService.revoke_email_code_login_token")
@patch("controllers.web.login.WebAppAuthService.get_email_code_login_data")
def test_should_normalize_email_before_validating(
self,
mock_get_token_data,
mock_revoke_token,
mock_get_user,
mock_login,
mock_reset_login_rate,
app,
):
mock_get_token_data.return_value = {"email": "User@Example.com", "code": "123456"}
mock_get_user.return_value = MagicMock()
with app.test_request_context(
"/web/email-code-login/validity",
method="POST",
json={"email": "User@Example.com", "code": encode_code("123456"), "token": "token-123"},
):
response = EmailCodeLoginApi().post()
assert response.get_json() == {"result": "success", "data": {"access_token": "new-access-token"}}
mock_get_user.assert_called_once_with("User@Example.com")
mock_revoke_token.assert_called_once_with("token-123")
mock_login.assert_called_once()
mock_reset_login_rate.assert_called_once_with("user@example.com")