mirror of
https://github.com/infiniflow/ragflow.git
synced 2026-07-19 14:11:04 +08:00
478c9846a1
### What problem does this PR solve? Remove tenant_llm call in rag. ### Type of change - [x] Bug Fix (non-breaking change which fixes an issue)
1620 lines
65 KiB
Python
1620 lines
65 KiB
Python
#
|
|
# Copyright 2026 The InfiniFlow Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
import asyncio
|
|
import base64
|
|
import importlib.util
|
|
import sys
|
|
from pathlib import Path
|
|
from types import ModuleType, SimpleNamespace
|
|
|
|
import pytest
|
|
|
|
|
|
class _DummyManager:
|
|
def route(self, *_args, **_kwargs):
|
|
def decorator(func):
|
|
return func
|
|
|
|
return decorator
|
|
|
|
|
|
class _AwaitableValue:
|
|
def __init__(self, value):
|
|
self._value = value
|
|
|
|
def __await__(self):
|
|
async def _co():
|
|
return self._value
|
|
|
|
return _co().__await__()
|
|
|
|
|
|
class _Field:
|
|
def __init__(self, name):
|
|
self.name = name
|
|
|
|
def __eq__(self, other):
|
|
return (self.name, other)
|
|
|
|
|
|
class _Invitee:
|
|
def __init__(self, user_id="invitee-1", email="invitee@example.com"):
|
|
self.id = user_id
|
|
self.email = email
|
|
|
|
def to_dict(self):
|
|
return {
|
|
"id": self.id,
|
|
"avatar": "avatar-url",
|
|
"email": self.email,
|
|
"nickname": "Invitee",
|
|
"password": "ignored",
|
|
}
|
|
|
|
|
|
def _run(coro):
|
|
return asyncio.run(coro)
|
|
|
|
|
|
def _passthrough_login_required(func):
|
|
async def _wrapper(*args, **kwargs):
|
|
return await func(*args, **kwargs)
|
|
|
|
_wrapper.__wrapped__ = func
|
|
return _wrapper
|
|
|
|
|
|
def _set_request_json(monkeypatch, module, payload):
|
|
async def _request_json():
|
|
return payload
|
|
|
|
monkeypatch.setattr(module, "get_request_json", _request_json)
|
|
|
|
|
|
def _load_tenant_module(monkeypatch):
|
|
repo_root = Path(__file__).resolve().parents[3]
|
|
|
|
api_pkg = ModuleType("api")
|
|
api_pkg.__path__ = [str(repo_root / "api")]
|
|
monkeypatch.setitem(sys.modules, "api", api_pkg)
|
|
|
|
apps_mod = ModuleType("api.apps")
|
|
apps_mod.__path__ = [str(repo_root / "api" / "apps")]
|
|
apps_mod.current_user = SimpleNamespace(id="tenant-1", email="owner@example.com")
|
|
apps_mod.login_required = lambda fn: fn
|
|
monkeypatch.setitem(sys.modules, "api.apps", apps_mod)
|
|
|
|
db_mod = ModuleType("api.db")
|
|
db_mod.UserTenantRole = SimpleNamespace(NORMAL="normal", OWNER="owner", INVITE="invite")
|
|
monkeypatch.setitem(sys.modules, "api.db", db_mod)
|
|
|
|
db_models_mod = ModuleType("api.db.db_models")
|
|
db_models_mod.UserTenant = type(
|
|
"UserTenant",
|
|
(),
|
|
{
|
|
"tenant_id": _Field("tenant_id"),
|
|
"user_id": _Field("user_id"),
|
|
},
|
|
)
|
|
monkeypatch.setitem(sys.modules, "api.db.db_models", db_models_mod)
|
|
|
|
services_pkg = ModuleType("api.db.services")
|
|
services_pkg.__path__ = []
|
|
monkeypatch.setitem(sys.modules, "api.db.services", services_pkg)
|
|
|
|
user_service_mod = ModuleType("api.db.services.user_service")
|
|
|
|
class _UserTenantService:
|
|
@staticmethod
|
|
def get_by_tenant_id(_tenant_id):
|
|
return []
|
|
|
|
@staticmethod
|
|
def query(**_kwargs):
|
|
return []
|
|
|
|
@staticmethod
|
|
def save(**_kwargs):
|
|
return True
|
|
|
|
@staticmethod
|
|
def filter_delete(_conditions):
|
|
return True
|
|
|
|
@staticmethod
|
|
def get_tenants_by_user_id(_user_id):
|
|
return []
|
|
|
|
@staticmethod
|
|
def filter_update(_conditions, _payload):
|
|
return True
|
|
|
|
class _UserService:
|
|
@staticmethod
|
|
def query(**_kwargs):
|
|
return []
|
|
|
|
@staticmethod
|
|
def get_by_id(_user_id):
|
|
return False, None
|
|
|
|
user_service_mod.UserTenantService = _UserTenantService
|
|
user_service_mod.UserService = _UserService
|
|
monkeypatch.setitem(sys.modules, "api.db.services.user_service", user_service_mod)
|
|
|
|
api_utils_mod = ModuleType("api.utils.api_utils")
|
|
api_utils_mod.get_json_result = lambda data=None, message="", code=0: {"code": code, "message": message, "data": data}
|
|
api_utils_mod.get_data_error_result = lambda message="": {"code": 102, "message": message, "data": False}
|
|
api_utils_mod.server_error_response = lambda exc: {"code": 100, "message": repr(exc), "data": False}
|
|
api_utils_mod.validate_request = lambda *_args, **_kwargs: (lambda fn: fn)
|
|
api_utils_mod.get_request_json = lambda: _AwaitableValue({})
|
|
monkeypatch.setitem(sys.modules, "api.utils.api_utils", api_utils_mod)
|
|
|
|
web_utils_mod = ModuleType("api.utils.web_utils")
|
|
web_utils_mod.send_invite_email = lambda **_kwargs: {"ok": True}
|
|
monkeypatch.setitem(sys.modules, "api.utils.web_utils", web_utils_mod)
|
|
|
|
common_pkg = ModuleType("common")
|
|
common_pkg.__path__ = [str(repo_root / "common")]
|
|
monkeypatch.setitem(sys.modules, "common", common_pkg)
|
|
|
|
constants_mod = ModuleType("common.constants")
|
|
constants_mod.RetCode = SimpleNamespace(AUTHENTICATION_ERROR=401, SERVER_ERROR=500, DATA_ERROR=102)
|
|
constants_mod.StatusEnum = SimpleNamespace(VALID=SimpleNamespace(value=1))
|
|
monkeypatch.setitem(sys.modules, "common.constants", constants_mod)
|
|
|
|
misc_utils_mod = ModuleType("common.misc_utils")
|
|
misc_utils_mod.get_uuid = lambda: "uuid-1"
|
|
monkeypatch.setitem(sys.modules, "common.misc_utils", misc_utils_mod)
|
|
|
|
time_utils_mod = ModuleType("common.time_utils")
|
|
time_utils_mod.delta_seconds = lambda _value: 0
|
|
monkeypatch.setitem(sys.modules, "common.time_utils", time_utils_mod)
|
|
|
|
settings_mod = ModuleType("common.settings")
|
|
settings_mod.MAIL_FRONTEND_URL = "https://frontend.example/invite"
|
|
monkeypatch.setitem(sys.modules, "common.settings", settings_mod)
|
|
common_pkg.settings = settings_mod
|
|
|
|
sys.modules.pop("test_tenant_app_unit_module", None)
|
|
module_path = repo_root / "api" / "apps" / "restful_apis" / "tenant_api.py"
|
|
spec = importlib.util.spec_from_file_location("test_tenant_app_unit_module", module_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
module.manager = _DummyManager()
|
|
monkeypatch.setitem(sys.modules, "test_tenant_app_unit_module", module)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
@pytest.mark.p2
|
|
def test_user_list_auth_success_exception_matrix_unit(monkeypatch):
|
|
module = _load_tenant_module(monkeypatch)
|
|
|
|
module.current_user.id = "other-user"
|
|
res = module.user_list("tenant-1")
|
|
assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res
|
|
assert res["message"] == "No authorization.", res
|
|
|
|
module.current_user.id = "tenant-1"
|
|
monkeypatch.setattr(
|
|
module.UserTenantService,
|
|
"get_by_tenant_id",
|
|
lambda _tenant_id: [{"id": "u1", "update_date": "2024-01-01 00:00:00"}],
|
|
)
|
|
monkeypatch.setattr(module, "delta_seconds", lambda _value: 42)
|
|
res = module.user_list("tenant-1")
|
|
assert res["code"] == 0, res
|
|
assert res["data"][0]["delta_seconds"] == 42, res
|
|
|
|
monkeypatch.setattr(module.UserTenantService, "get_by_tenant_id", lambda _tenant_id: (_ for _ in ()).throw(RuntimeError("list boom")))
|
|
res = module.user_list("tenant-1")
|
|
assert res["code"] == 100, res
|
|
assert "list boom" in res["message"], res
|
|
|
|
|
|
@pytest.mark.p2
|
|
def test_create_invite_role_and_email_failure_matrix_unit(monkeypatch):
|
|
module = _load_tenant_module(monkeypatch)
|
|
|
|
module.current_user.id = "other-user"
|
|
_set_request_json(monkeypatch, module, {"email": "invitee@example.com"})
|
|
res = _run(module.create("tenant-1"))
|
|
assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res
|
|
assert res["message"] == "No authorization.", res
|
|
|
|
module.current_user.id = "tenant-1"
|
|
monkeypatch.setattr(module.UserService, "query", lambda **_kwargs: [])
|
|
res = _run(module.create("tenant-1"))
|
|
assert res["message"] == "User not found.", res
|
|
|
|
invitee = _Invitee()
|
|
monkeypatch.setattr(module.UserService, "query", lambda **_kwargs: [invitee])
|
|
monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(role=module.UserTenantRole.NORMAL)])
|
|
res = _run(module.create("tenant-1"))
|
|
assert "already in the team." in res["message"], res
|
|
|
|
monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(role=module.UserTenantRole.OWNER)])
|
|
res = _run(module.create("tenant-1"))
|
|
assert "owner of the team." in res["message"], res
|
|
|
|
monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(role="strange-role")])
|
|
res = _run(module.create("tenant-1"))
|
|
assert "role: strange-role is invalid." in res["message"], res
|
|
|
|
saved = []
|
|
scheduled = []
|
|
monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [])
|
|
monkeypatch.setattr(module.UserTenantService, "save", lambda **kwargs: saved.append(kwargs) or True)
|
|
monkeypatch.setattr(module.UserService, "get_by_id", lambda _user_id: (True, SimpleNamespace(nickname="Inviter Nick")))
|
|
monkeypatch.setattr(module, "send_invite_email", lambda **kwargs: kwargs)
|
|
monkeypatch.setattr(module.asyncio, "create_task", lambda payload: scheduled.append(payload) or SimpleNamespace())
|
|
res = _run(module.create("tenant-1"))
|
|
assert res["code"] == 0, res
|
|
assert saved and saved[-1]["role"] == module.UserTenantRole.INVITE, saved
|
|
assert scheduled and scheduled[-1]["inviter"] == "Inviter Nick", scheduled
|
|
assert sorted(res["data"].keys()) == ["avatar", "email", "id", "nickname"], res
|
|
|
|
monkeypatch.setattr(module.asyncio, "create_task", lambda _payload: (_ for _ in ()).throw(RuntimeError("send boom")))
|
|
res = _run(module.create("tenant-1"))
|
|
assert res["code"] == module.RetCode.SERVER_ERROR, res
|
|
assert "Failed to send invite email." in res["message"], res
|
|
|
|
|
|
@pytest.mark.p2
|
|
def test_rm_and_tenant_list_matrix_unit(monkeypatch):
|
|
module = _load_tenant_module(monkeypatch)
|
|
|
|
module.current_user.id = "outsider"
|
|
_set_request_json(monkeypatch, module, {"user_id": "user-2"})
|
|
res = _run(module.rm("tenant-1"))
|
|
assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res
|
|
assert res["message"] == "No authorization.", res
|
|
|
|
module.current_user.id = "tenant-1"
|
|
deleted = []
|
|
monkeypatch.setattr(module.UserTenantService, "filter_delete", lambda conditions: deleted.append(conditions) or True)
|
|
res = _run(module.rm("tenant-1"))
|
|
assert res["code"] == 0, res
|
|
assert res["data"] is True, res
|
|
assert deleted, "filter_delete should be called"
|
|
|
|
monkeypatch.setattr(module.UserTenantService, "filter_delete", lambda _conditions: (_ for _ in ()).throw(RuntimeError("rm boom")))
|
|
res = _run(module.rm("tenant-1"))
|
|
assert res["code"] == 100, res
|
|
assert "rm boom" in res["message"], res
|
|
|
|
monkeypatch.setattr(
|
|
module.UserTenantService,
|
|
"get_tenants_by_user_id",
|
|
lambda _user_id: [{"id": "tenant-1", "update_date": "2024-01-01 00:00:00"}],
|
|
)
|
|
monkeypatch.setattr(module, "delta_seconds", lambda _value: 9)
|
|
res = module.tenant_list()
|
|
assert res["code"] == 0, res
|
|
assert res["data"][0]["delta_seconds"] == 9, res
|
|
|
|
monkeypatch.setattr(module.UserTenantService, "get_tenants_by_user_id", lambda _user_id: (_ for _ in ()).throw(RuntimeError("tenant boom")))
|
|
res = module.tenant_list()
|
|
assert res["code"] == 100, res
|
|
assert "tenant boom" in res["message"], res
|
|
|
|
|
|
@pytest.mark.p2
|
|
def test_agree_success_and_exception_unit(monkeypatch):
|
|
module = _load_tenant_module(monkeypatch)
|
|
|
|
calls = []
|
|
monkeypatch.setattr(module.UserTenantService, "filter_update", lambda conditions, payload: calls.append((conditions, payload)) or True)
|
|
res = module.agree("tenant-1")
|
|
assert res["code"] == 0, res
|
|
assert res["data"] is True, res
|
|
assert calls and calls[-1][1]["role"] == module.UserTenantRole.NORMAL
|
|
|
|
monkeypatch.setattr(module.UserTenantService, "filter_update", lambda _conditions, _payload: (_ for _ in ()).throw(RuntimeError("agree boom")))
|
|
res = module.agree("tenant-1")
|
|
assert res["code"] == 100, res
|
|
assert "agree boom" in res["message"], res
|
|
|
|
|
|
class _Args(dict):
|
|
def get(self, key, default=None, type=None):
|
|
value = super().get(key, default)
|
|
if type is None:
|
|
return value
|
|
try:
|
|
return type(value)
|
|
except (TypeError, ValueError):
|
|
return default
|
|
|
|
|
|
class _DummyResponse:
|
|
def __init__(self, data):
|
|
self.data = data
|
|
self.headers = {}
|
|
|
|
|
|
class _DummyHTTPResponse:
|
|
def __init__(self, payload):
|
|
self._payload = payload
|
|
|
|
def json(self):
|
|
return self._payload
|
|
|
|
|
|
class _DummyRedis:
|
|
def __init__(self):
|
|
self.store = {}
|
|
|
|
def get(self, key):
|
|
return self.store.get(key)
|
|
|
|
def set(self, key, value, _ttl=None):
|
|
self.store[key] = value
|
|
|
|
def delete(self, key):
|
|
self.store.pop(key, None)
|
|
|
|
|
|
class _DummyUser:
|
|
def __init__(self, user_id, email, *, password="stored-password", is_active="1", nickname="nick"):
|
|
self.id = user_id
|
|
self.email = email
|
|
self.password = password
|
|
self.is_active = is_active
|
|
self.nickname = nickname
|
|
self.access_token = ""
|
|
self.save_calls = 0
|
|
|
|
def save(self):
|
|
self.save_calls += 1
|
|
|
|
def get_id(self):
|
|
return self.id
|
|
|
|
def to_json(self):
|
|
return {"id": self.id, "email": self.email, "nickname": self.nickname}
|
|
|
|
def to_dict(self):
|
|
return {"id": self.id, "email": self.email}
|
|
|
|
def to_safe_dict(self, *, for_self: bool = False):
|
|
_sensitive = {"password", "access_token", "email"}
|
|
result = {k: v for k, v in self.to_dict().items() if k not in _sensitive}
|
|
if for_self:
|
|
result["email"] = self.email
|
|
return result
|
|
|
|
|
|
def _set_request_args(monkeypatch, module, args=None):
|
|
monkeypatch.setattr(module, "request", SimpleNamespace(args=_Args(args or {})))
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def auth():
|
|
return "unit-auth"
|
|
|
|
|
|
@pytest.fixture(scope="session", autouse=True)
|
|
def set_tenant_info():
|
|
return None
|
|
|
|
|
|
def _load_user_app(monkeypatch):
|
|
repo_root = Path(__file__).resolve().parents[3]
|
|
|
|
quart_mod = ModuleType("quart")
|
|
quart_mod.session = {}
|
|
quart_mod.request = SimpleNamespace(args=_Args({}))
|
|
|
|
async def _make_response(data):
|
|
return _DummyResponse(data)
|
|
|
|
quart_mod.make_response = _make_response
|
|
quart_mod.redirect = lambda url: {"redirect": url}
|
|
monkeypatch.setitem(sys.modules, "quart", quart_mod)
|
|
|
|
api_pkg = ModuleType("api")
|
|
api_pkg.__path__ = [str(repo_root / "api")]
|
|
monkeypatch.setitem(sys.modules, "api", api_pkg)
|
|
|
|
apps_mod = ModuleType("api.apps")
|
|
apps_mod.__path__ = [str(repo_root / "api" / "apps")]
|
|
apps_mod.current_user = _DummyUser("current-user", "current@example.com")
|
|
apps_mod.login_required = lambda fn: fn
|
|
apps_mod.login_user = lambda _user: True
|
|
apps_mod.logout_user = lambda: True
|
|
monkeypatch.setitem(sys.modules, "api.apps", apps_mod)
|
|
api_pkg.apps = apps_mod
|
|
|
|
apps_auth_mod = ModuleType("api.apps.auth")
|
|
apps_auth_mod.get_auth_client = lambda _config: SimpleNamespace(
|
|
get_authorization_url=lambda state: f"https://oauth.example/{state}"
|
|
)
|
|
monkeypatch.setitem(sys.modules, "api.apps.auth", apps_auth_mod)
|
|
|
|
db_mod = ModuleType("api.db")
|
|
db_mod.FileType = SimpleNamespace(FOLDER=SimpleNamespace(value="folder"))
|
|
db_mod.UserTenantRole = SimpleNamespace(OWNER="owner")
|
|
monkeypatch.setitem(sys.modules, "api.db", db_mod)
|
|
api_pkg.db = db_mod
|
|
|
|
db_models_mod = ModuleType("api.db.db_models")
|
|
|
|
class _DummyTenantLLMModel:
|
|
tenant_id = _Field("tenant_id")
|
|
|
|
@staticmethod
|
|
def delete():
|
|
class _DeleteQuery:
|
|
def where(self, *_args, **_kwargs):
|
|
return self
|
|
|
|
def execute(self):
|
|
return 1
|
|
|
|
return _DeleteQuery()
|
|
|
|
db_models_mod.TenantLLM = _DummyTenantLLMModel
|
|
monkeypatch.setitem(sys.modules, "api.db.db_models", db_models_mod)
|
|
|
|
services_pkg = ModuleType("api.db.services")
|
|
services_pkg.__path__ = []
|
|
monkeypatch.setitem(sys.modules, "api.db.services", services_pkg)
|
|
|
|
file_service_mod = ModuleType("api.db.services.file_service")
|
|
|
|
class _StubFileService:
|
|
@staticmethod
|
|
def insert(_data):
|
|
return True
|
|
|
|
file_service_mod.FileService = _StubFileService
|
|
monkeypatch.setitem(sys.modules, "api.db.services.file_service", file_service_mod)
|
|
|
|
llm_service_mod = ModuleType("api.db.services.llm_service")
|
|
llm_service_mod.get_init_tenant_llm = lambda _user_id: []
|
|
monkeypatch.setitem(sys.modules, "api.db.services.llm_service", llm_service_mod)
|
|
|
|
tenant_llm_service_mod = ModuleType("api.db.services.tenant_llm_service")
|
|
|
|
class _MockTableObject:
|
|
def __init__(self, **kwargs):
|
|
for key, value in kwargs.items():
|
|
setattr(self, key, value)
|
|
|
|
def to_dict(self):
|
|
return {k: v for k, v in self.__dict__.items()}
|
|
|
|
class _StubTenantLLMService:
|
|
@staticmethod
|
|
def insert_many(_payload):
|
|
return True
|
|
|
|
@staticmethod
|
|
def get_api_key(tenant_id, model_name, model_type=None):
|
|
return _MockTableObject(
|
|
id=1,
|
|
tenant_id=tenant_id,
|
|
llm_factory="",
|
|
model_type="chat",
|
|
llm_name=model_name,
|
|
api_key="fake-api-key",
|
|
api_base="https://api.example.com",
|
|
max_tokens=8192,
|
|
used_tokens=0,
|
|
status=1
|
|
)
|
|
|
|
tenant_llm_service_mod.TenantLLMService = _StubTenantLLMService
|
|
monkeypatch.setitem(sys.modules, "api.db.services.tenant_llm_service", tenant_llm_service_mod)
|
|
|
|
user_service_mod = ModuleType("api.db.services.user_service")
|
|
|
|
class _StubTenantService:
|
|
@staticmethod
|
|
def insert(**_kwargs):
|
|
return True
|
|
|
|
@staticmethod
|
|
def delete_by_id(_tenant_id):
|
|
return True
|
|
|
|
@staticmethod
|
|
def get_by_id(_tenant_id):
|
|
return True, SimpleNamespace(id=_tenant_id)
|
|
|
|
@staticmethod
|
|
def get_info_by(_user_id):
|
|
return []
|
|
|
|
@staticmethod
|
|
def update_by_id(_tenant_id, _payload):
|
|
return True
|
|
|
|
class _StubUserService:
|
|
@staticmethod
|
|
def query(**_kwargs):
|
|
return []
|
|
|
|
@staticmethod
|
|
def query_user(_email, _password):
|
|
return None
|
|
|
|
@staticmethod
|
|
def query_user_by_email(**_kwargs):
|
|
return []
|
|
|
|
@staticmethod
|
|
def save(**_kwargs):
|
|
return True
|
|
|
|
@staticmethod
|
|
def delete_by_id(_user_id):
|
|
return True
|
|
|
|
@staticmethod
|
|
def update_by_id(_user_id, _payload):
|
|
return True
|
|
|
|
@staticmethod
|
|
def update_user_password(_user_id, _new_password):
|
|
return True
|
|
|
|
class _StubUserTenantService:
|
|
@staticmethod
|
|
def insert(**_kwargs):
|
|
return True
|
|
|
|
@staticmethod
|
|
def query(**_kwargs):
|
|
return []
|
|
|
|
@staticmethod
|
|
def delete_by_id(_user_tenant_id):
|
|
return True
|
|
|
|
user_service_mod.TenantService = _StubTenantService
|
|
user_service_mod.UserService = _StubUserService
|
|
user_service_mod.UserTenantService = _StubUserTenantService
|
|
monkeypatch.setitem(sys.modules, "api.db.services.user_service", user_service_mod)
|
|
|
|
api_utils_mod = ModuleType("api.utils.api_utils")
|
|
|
|
async def _default_request_json():
|
|
return {}
|
|
|
|
def _get_json_result(code=0, message="success", data=None):
|
|
return {"code": code, "message": message, "data": data}
|
|
|
|
def _get_data_error_result(code=102, message="Sorry! Data missing!", data=None):
|
|
return {"code": code, "message": message, "data": data}
|
|
|
|
def _server_error_response(error):
|
|
return {"code": 100, "message": repr(error)}
|
|
|
|
def _validate_request(*_args, **_kwargs):
|
|
def _decorator(func):
|
|
return func
|
|
|
|
return _decorator
|
|
|
|
api_utils_mod.get_request_json = _default_request_json
|
|
api_utils_mod.get_json_result = _get_json_result
|
|
api_utils_mod.get_data_error_result = _get_data_error_result
|
|
api_utils_mod.server_error_response = _server_error_response
|
|
api_utils_mod.validate_request = _validate_request
|
|
monkeypatch.setitem(sys.modules, "api.utils.api_utils", api_utils_mod)
|
|
|
|
crypt_mod = ModuleType("api.utils.crypt")
|
|
crypt_mod.decrypt = lambda value: value
|
|
monkeypatch.setitem(sys.modules, "api.utils.crypt", crypt_mod)
|
|
|
|
web_utils_mod = ModuleType("api.utils.web_utils")
|
|
web_utils_mod.send_email_html = lambda *_args, **_kwargs: _AwaitableValue(True)
|
|
web_utils_mod.OTP_LENGTH = 6
|
|
web_utils_mod.OTP_TTL_SECONDS = 600
|
|
web_utils_mod.ATTEMPT_LIMIT = 5
|
|
web_utils_mod.ATTEMPT_LOCK_SECONDS = 600
|
|
web_utils_mod.RESEND_COOLDOWN_SECONDS = 60
|
|
web_utils_mod.otp_keys = lambda email: (
|
|
f"otp:{email}:code",
|
|
f"otp:{email}:attempts",
|
|
f"otp:{email}:last",
|
|
f"otp:{email}:lock",
|
|
)
|
|
web_utils_mod.hash_code = lambda code, _salt: f"hash:{code}"
|
|
web_utils_mod.captcha_key = lambda email: f"captcha:{email}"
|
|
monkeypatch.setitem(sys.modules, "api.utils.web_utils", web_utils_mod)
|
|
|
|
common_pkg = ModuleType("common")
|
|
common_pkg.__path__ = [str(repo_root / "common")]
|
|
monkeypatch.setitem(sys.modules, "common", common_pkg)
|
|
|
|
settings_mod = ModuleType("common.settings")
|
|
settings_mod.OAUTH_CONFIG = {
|
|
"github": {"display_name": "GitHub", "icon": "gh"},
|
|
"feishu": {"display_name": "Feishu", "icon": "fs"},
|
|
}
|
|
settings_mod.GITHUB_OAUTH = {"url": "https://github.example/oauth", "client_id": "cid", "secret_key": "sk"}
|
|
settings_mod.FEISHU_OAUTH = {
|
|
"app_access_token_url": "https://feishu.example/app_token",
|
|
"user_access_token_url": "https://feishu.example/user_token",
|
|
"app_id": "app-id",
|
|
"app_secret": "app-secret",
|
|
"grant_type": "authorization_code",
|
|
}
|
|
settings_mod.CHAT_MDL = "chat-mdl"
|
|
settings_mod.EMBEDDING_MDL = "embd-mdl"
|
|
settings_mod.ASR_MDL = "asr-mdl"
|
|
settings_mod.PARSERS = []
|
|
settings_mod.IMAGE2TEXT_MDL = "img-mdl"
|
|
settings_mod.RERANK_MDL = "rerank-mdl"
|
|
settings_mod.REGISTER_ENABLED = True
|
|
monkeypatch.setitem(sys.modules, "common.settings", settings_mod)
|
|
common_pkg.settings = settings_mod
|
|
|
|
constants_mod = ModuleType("common.constants")
|
|
constants_mod.RetCode = SimpleNamespace(
|
|
AUTHENTICATION_ERROR=401,
|
|
SERVER_ERROR=500,
|
|
FORBIDDEN=403,
|
|
EXCEPTION_ERROR=100,
|
|
OPERATING_ERROR=300,
|
|
ARGUMENT_ERROR=101,
|
|
DATA_ERROR=102,
|
|
NOT_EFFECTIVE=103,
|
|
SUCCESS=0,
|
|
)
|
|
monkeypatch.setitem(sys.modules, "common.constants", constants_mod)
|
|
|
|
connection_utils_mod = ModuleType("common.connection_utils")
|
|
|
|
async def _construct_response(data=None, auth=None, message=""):
|
|
return {"code": 0, "message": message, "data": data, "auth": auth}
|
|
|
|
connection_utils_mod.construct_response = _construct_response
|
|
monkeypatch.setitem(sys.modules, "common.connection_utils", connection_utils_mod)
|
|
|
|
time_utils_mod = ModuleType("common.time_utils")
|
|
time_utils_mod.current_timestamp = lambda: 111
|
|
time_utils_mod.datetime_format = lambda _dt: "2024-01-01 00:00:00"
|
|
time_utils_mod.get_format_time = lambda: "2024-01-01 00:00:00"
|
|
monkeypatch.setitem(sys.modules, "common.time_utils", time_utils_mod)
|
|
|
|
misc_utils_mod = ModuleType("common.misc_utils")
|
|
misc_utils_mod.download_img = lambda _url: "avatar"
|
|
misc_utils_mod.get_uuid = lambda: "uuid-default"
|
|
monkeypatch.setitem(sys.modules, "common.misc_utils", misc_utils_mod)
|
|
|
|
http_client_mod = ModuleType("common.http_client")
|
|
|
|
async def _async_request(_method, _url, **_kwargs):
|
|
return _DummyHTTPResponse({})
|
|
|
|
http_client_mod.async_request = _async_request
|
|
monkeypatch.setitem(sys.modules, "common.http_client", http_client_mod)
|
|
|
|
rag_pkg = ModuleType("rag")
|
|
rag_pkg.__path__ = [str(repo_root / "rag")]
|
|
monkeypatch.setitem(sys.modules, "rag", rag_pkg)
|
|
|
|
rag_utils_pkg = ModuleType("rag.utils")
|
|
rag_utils_pkg.__path__ = [str(repo_root / "rag" / "utils")]
|
|
monkeypatch.setitem(sys.modules, "rag.utils", rag_utils_pkg)
|
|
|
|
redis_mod = ModuleType("rag.utils.redis_conn")
|
|
redis_mod.REDIS_CONN = _DummyRedis()
|
|
monkeypatch.setitem(sys.modules, "rag.utils.redis_conn", redis_mod)
|
|
|
|
module_name = "test_user_app_unit_module"
|
|
module_path = repo_root / "api" / "apps" / "restful_apis" / "user_api.py"
|
|
spec = importlib.util.spec_from_file_location(module_name, module_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
module.manager = _DummyManager()
|
|
monkeypatch.setitem(sys.modules, module_name, module)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
@pytest.mark.p2
|
|
def test_login_route_branch_matrix_unit(monkeypatch):
|
|
module = _load_user_app(monkeypatch)
|
|
|
|
_set_request_json(monkeypatch, module, {})
|
|
res = _run(module.login())
|
|
assert res["code"] == module.RetCode.AUTHENTICATION_ERROR
|
|
assert "Unauthorized" in res["message"]
|
|
|
|
_set_request_json(monkeypatch, module, {"email": "unknown@example.com", "password": "enc"})
|
|
monkeypatch.setattr(module.UserService, "query", lambda **_kwargs: [])
|
|
res = _run(module.login())
|
|
assert res["code"] == module.RetCode.AUTHENTICATION_ERROR
|
|
assert "not registered" in res["message"]
|
|
|
|
_set_request_json(monkeypatch, module, {"email": "known@example.com", "password": "enc"})
|
|
monkeypatch.setattr(module.UserService, "query", lambda **_kwargs: [SimpleNamespace(email="known@example.com")])
|
|
|
|
def _raise_decrypt(_value):
|
|
raise RuntimeError("decrypt explode")
|
|
|
|
monkeypatch.setattr(module, "decrypt", _raise_decrypt)
|
|
res = _run(module.login())
|
|
assert res["code"] == module.RetCode.SERVER_ERROR
|
|
assert "Fail to crypt password" in res["message"]
|
|
|
|
user_inactive = _DummyUser("u-inactive", "known@example.com", is_active="0")
|
|
monkeypatch.setattr(module, "decrypt", lambda value: value)
|
|
monkeypatch.setattr(module.UserService, "query_user", lambda _email, _password: user_inactive)
|
|
res = _run(module.login())
|
|
assert res["code"] == module.RetCode.FORBIDDEN
|
|
assert "disabled" in res["message"]
|
|
|
|
monkeypatch.setattr(module.UserService, "query_user", lambda _email, _password: None)
|
|
res = _run(module.login())
|
|
assert res["code"] == module.RetCode.AUTHENTICATION_ERROR
|
|
assert "do not match" in res["message"]
|
|
|
|
|
|
@pytest.mark.p2
|
|
def test_login_channels_and_oauth_login_matrix_unit(monkeypatch):
|
|
module = _load_user_app(monkeypatch)
|
|
|
|
module.settings.OAUTH_CONFIG = {"github": {"display_name": "GitHub", "icon": "gh"}}
|
|
res = _run(module.get_login_channels())
|
|
assert res["code"] == 0
|
|
assert res["data"][0]["channel"] == "github"
|
|
|
|
class _BrokenOAuthConfig:
|
|
@staticmethod
|
|
def items():
|
|
raise RuntimeError("broken oauth config")
|
|
|
|
module.settings.OAUTH_CONFIG = _BrokenOAuthConfig()
|
|
res = _run(module.get_login_channels())
|
|
assert res["code"] == module.RetCode.EXCEPTION_ERROR
|
|
assert "Load channels failure" in res["message"]
|
|
|
|
module.settings.OAUTH_CONFIG = {"github": {"display_name": "GitHub", "icon": "gh"}}
|
|
with pytest.raises(ValueError, match="Invalid channel name: missing"):
|
|
_run(module.oauth_login("missing"))
|
|
|
|
module.session.clear()
|
|
monkeypatch.setattr(module, "get_uuid", lambda: "state-123")
|
|
|
|
class _AuthClient:
|
|
@staticmethod
|
|
def get_authorization_url(state):
|
|
return f"https://oauth.example/{state}"
|
|
|
|
monkeypatch.setattr(module, "get_auth_client", lambda _config: _AuthClient())
|
|
res = _run(module.oauth_login("github"))
|
|
assert res["redirect"] == "https://oauth.example/state-123"
|
|
assert module.session["oauth_state"] == "state-123"
|
|
|
|
|
|
@pytest.mark.p2
|
|
def test_oauth_callback_matrix_unit(monkeypatch):
|
|
module = _load_user_app(monkeypatch)
|
|
module.settings.OAUTH_CONFIG = {"github": {"display_name": "GitHub", "icon": "gh"}}
|
|
|
|
class _SyncAuthClient:
|
|
def __init__(self, token_info, user_info):
|
|
self._token_info = token_info
|
|
self._user_info = user_info
|
|
|
|
def exchange_code_for_token(self, _code):
|
|
return self._token_info
|
|
|
|
def fetch_user_info(self, _token, id_token=None):
|
|
_ = id_token
|
|
return self._user_info
|
|
|
|
class _AsyncAuthClient:
|
|
def __init__(self, token_info, user_info):
|
|
self._token_info = token_info
|
|
self._user_info = user_info
|
|
|
|
async def async_exchange_code_for_token(self, _code):
|
|
return self._token_info
|
|
|
|
async def async_fetch_user_info(self, _token, id_token=None):
|
|
_ = id_token
|
|
return self._user_info
|
|
|
|
_set_request_args(monkeypatch, module, {"state": "x", "code": "c"})
|
|
module.session.clear()
|
|
res = _run(module.oauth_callback("missing"))
|
|
assert "Invalid channel name: missing" in res["redirect"]
|
|
|
|
sync_ok = _SyncAuthClient(
|
|
token_info={"access_token": "token-sync", "id_token": "id-sync"},
|
|
user_info=SimpleNamespace(email="sync@example.com", avatar_url="http://img", nickname="sync"),
|
|
)
|
|
monkeypatch.setattr(module, "get_auth_client", lambda _config: sync_ok)
|
|
|
|
module.session.clear()
|
|
module.session["oauth_state"] = "expected"
|
|
_set_request_args(monkeypatch, module, {"state": "wrong", "code": "code"})
|
|
res = _run(module.oauth_callback("github"))
|
|
assert res["redirect"] == "/?error=invalid_state"
|
|
|
|
module.session.clear()
|
|
module.session["oauth_state"] = "ok-state"
|
|
_set_request_args(monkeypatch, module, {"state": "ok-state"})
|
|
res = _run(module.oauth_callback("github"))
|
|
assert res["redirect"] == "/?error=missing_code"
|
|
|
|
sync_missing_token = _SyncAuthClient(
|
|
token_info={"id_token": "id-only"},
|
|
user_info=SimpleNamespace(email="sync@example.com", avatar_url="http://img", nickname="sync"),
|
|
)
|
|
monkeypatch.setattr(module, "get_auth_client", lambda _config: sync_missing_token)
|
|
module.session.clear()
|
|
module.session["oauth_state"] = "token-state"
|
|
_set_request_args(monkeypatch, module, {"state": "token-state", "code": "code"})
|
|
res = _run(module.oauth_callback("github"))
|
|
assert res["redirect"] == "/?error=token_failed"
|
|
|
|
sync_missing_email = _SyncAuthClient(
|
|
token_info={"access_token": "token-sync", "id_token": "id-sync"},
|
|
user_info=SimpleNamespace(email=None, avatar_url="http://img", nickname="sync"),
|
|
)
|
|
monkeypatch.setattr(module, "get_auth_client", lambda _config: sync_missing_email)
|
|
module.session.clear()
|
|
module.session["oauth_state"] = "email-state"
|
|
_set_request_args(monkeypatch, module, {"state": "email-state", "code": "code"})
|
|
res = _run(module.oauth_callback("github"))
|
|
assert res["redirect"] == "/?error=email_missing"
|
|
|
|
async_new_user = _AsyncAuthClient(
|
|
token_info={"access_token": "token-async", "id_token": "id-async"},
|
|
user_info=SimpleNamespace(email="new@example.com", avatar_url="http://img", nickname="new-user"),
|
|
)
|
|
monkeypatch.setattr(module, "get_auth_client", lambda _config: async_new_user)
|
|
monkeypatch.setattr(module.UserService, "query", lambda **_kwargs: [])
|
|
|
|
def _raise_download(_url):
|
|
raise RuntimeError("download explode")
|
|
|
|
monkeypatch.setattr(module, "download_img", _raise_download)
|
|
monkeypatch.setattr(module, "user_register", lambda _user_id, _user: None)
|
|
rollback_calls = []
|
|
monkeypatch.setattr(module, "rollback_user_registration", lambda user_id: rollback_calls.append(user_id))
|
|
monkeypatch.setattr(module, "get_uuid", lambda: "new-user-id")
|
|
module.session.clear()
|
|
module.session["oauth_state"] = "new-user-state"
|
|
_set_request_args(monkeypatch, module, {"state": "new-user-state", "code": "code"})
|
|
res = _run(module.oauth_callback("github"))
|
|
assert "Failed to register new@example.com" in res["redirect"]
|
|
assert rollback_calls == ["new-user-id"]
|
|
|
|
monkeypatch.setattr(module, "download_img", lambda _url: "avatar")
|
|
monkeypatch.setattr(
|
|
module,
|
|
"user_register",
|
|
lambda _user_id, _user: [_DummyUser("dup-1", "new@example.com"), _DummyUser("dup-2", "new@example.com")],
|
|
)
|
|
rollback_calls.clear()
|
|
module.session.clear()
|
|
module.session["oauth_state"] = "dup-user-state"
|
|
_set_request_args(monkeypatch, module, {"state": "dup-user-state", "code": "code"})
|
|
res = _run(module.oauth_callback("github"))
|
|
assert "Same email: new@example.com exists!" in res["redirect"]
|
|
assert rollback_calls == ["new-user-id"]
|
|
|
|
new_user = _DummyUser("new-user", "new@example.com")
|
|
login_calls = []
|
|
monkeypatch.setattr(module, "login_user", lambda user: login_calls.append(user))
|
|
monkeypatch.setattr(module, "user_register", lambda _user_id, _user: [new_user])
|
|
module.session.clear()
|
|
module.session["oauth_state"] = "create-user-state"
|
|
_set_request_args(monkeypatch, module, {"state": "create-user-state", "code": "code"})
|
|
res = _run(module.oauth_callback("github"))
|
|
assert res["redirect"] == "/?auth=new-user"
|
|
assert login_calls and login_calls[-1] is new_user
|
|
|
|
async_existing_inactive = _AsyncAuthClient(
|
|
token_info={"access_token": "token-existing", "id_token": "id-existing"},
|
|
user_info=SimpleNamespace(email="existing@example.com", avatar_url="http://img", nickname="existing"),
|
|
)
|
|
monkeypatch.setattr(module, "get_auth_client", lambda _config: async_existing_inactive)
|
|
inactive_user = _DummyUser("existing-user", "existing@example.com", is_active="0")
|
|
monkeypatch.setattr(module.UserService, "query", lambda **_kwargs: [inactive_user])
|
|
module.session.clear()
|
|
module.session["oauth_state"] = "inactive-state"
|
|
_set_request_args(monkeypatch, module, {"state": "inactive-state", "code": "code"})
|
|
res = _run(module.oauth_callback("github"))
|
|
assert res["redirect"] == "/?error=user_inactive"
|
|
|
|
async_existing_ok = _AsyncAuthClient(
|
|
token_info={"access_token": "token-existing", "id_token": "id-existing"},
|
|
user_info=SimpleNamespace(email="existing@example.com", avatar_url="http://img", nickname="existing"),
|
|
)
|
|
monkeypatch.setattr(module, "get_auth_client", lambda _config: async_existing_ok)
|
|
existing_user = _DummyUser("existing-user", "existing@example.com")
|
|
monkeypatch.setattr(module.UserService, "query", lambda **_kwargs: [existing_user])
|
|
login_calls.clear()
|
|
monkeypatch.setattr(module, "login_user", lambda user: login_calls.append(user))
|
|
monkeypatch.setattr(module, "get_uuid", lambda: "existing-token")
|
|
module.session.clear()
|
|
module.session["oauth_state"] = "existing-state"
|
|
_set_request_args(monkeypatch, module, {"state": "existing-state", "code": "code"})
|
|
res = _run(module.oauth_callback("github"))
|
|
assert res["redirect"] == "/?auth=existing-user"
|
|
assert existing_user.access_token == "existing-token"
|
|
assert existing_user.save_calls == 1
|
|
assert login_calls and login_calls[-1] is existing_user
|
|
|
|
|
|
@pytest.mark.p2
|
|
def test_logout_setting_profile_matrix_unit(monkeypatch):
|
|
module = _load_user_app(monkeypatch)
|
|
|
|
current_user = _DummyUser("current-user", "current@example.com", password="stored-password")
|
|
monkeypatch.setattr(module, "current_user", current_user)
|
|
monkeypatch.setattr(module.secrets, "token_hex", lambda _n: "abcdef")
|
|
logout_calls = []
|
|
monkeypatch.setattr(module, "logout_user", lambda: logout_calls.append(True))
|
|
|
|
res = _run(module.log_out())
|
|
assert res["code"] == 0
|
|
assert current_user.access_token == "INVALID_abcdef"
|
|
assert current_user.save_calls == 1
|
|
assert logout_calls == [True]
|
|
|
|
_set_request_json(monkeypatch, module, {"password": "old-password", "new_password": "new-password"})
|
|
monkeypatch.setattr(module, "decrypt", lambda value: value)
|
|
monkeypatch.setattr(module, "check_password_hash", lambda _hashed, _plain: False)
|
|
res = _run(module.setting_user())
|
|
assert res["code"] == module.RetCode.AUTHENTICATION_ERROR
|
|
assert "Password error" in res["message"]
|
|
|
|
_set_request_json(
|
|
monkeypatch,
|
|
module,
|
|
{
|
|
"password": "old-password",
|
|
"new_password": "new-password",
|
|
"nickname": "neo",
|
|
"email": "blocked@example.com",
|
|
"status": "disabled",
|
|
"theme": "dark",
|
|
},
|
|
)
|
|
monkeypatch.setattr(module, "check_password_hash", lambda _hashed, _plain: True)
|
|
monkeypatch.setattr(module, "decrypt", lambda value: f"dec:{value}")
|
|
monkeypatch.setattr(module, "generate_password_hash", lambda value: f"hash:{value}")
|
|
update_calls = {}
|
|
|
|
def _update_by_id(user_id, payload):
|
|
update_calls["user_id"] = user_id
|
|
update_calls["payload"] = payload
|
|
return True
|
|
|
|
monkeypatch.setattr(module.UserService, "update_by_id", _update_by_id)
|
|
res = _run(module.setting_user())
|
|
assert res["code"] == 0
|
|
assert res["data"] is True
|
|
assert update_calls["user_id"] == "current-user"
|
|
assert update_calls["payload"]["password"] == "hash:dec:new-password"
|
|
assert update_calls["payload"]["nickname"] == "neo"
|
|
assert update_calls["payload"]["theme"] == "dark"
|
|
assert "email" not in update_calls["payload"]
|
|
assert "status" not in update_calls["payload"]
|
|
|
|
_set_request_json(monkeypatch, module, {"nickname": "neo"})
|
|
|
|
def _raise_update(_user_id, _payload):
|
|
raise RuntimeError("update explode")
|
|
|
|
monkeypatch.setattr(module.UserService, "update_by_id", _raise_update)
|
|
res = _run(module.setting_user())
|
|
assert res["code"] == module.RetCode.EXCEPTION_ERROR
|
|
assert "Update failure" in res["message"]
|
|
|
|
res = _run(module.user_profile())
|
|
assert res["code"] == 0
|
|
assert res["data"] == current_user.to_dict()
|
|
|
|
|
|
@pytest.mark.p2
|
|
def test_registration_helpers_and_register_route_matrix_unit(monkeypatch):
|
|
module = _load_user_app(monkeypatch)
|
|
|
|
deleted = {"user": 0, "tenant": 0, "user_tenant": 0}
|
|
monkeypatch.setattr(module.UserService, "delete_by_id", lambda _user_id: deleted.__setitem__("user", deleted["user"] + 1))
|
|
monkeypatch.setattr(module.TenantService, "delete_by_id", lambda _tenant_id: deleted.__setitem__("tenant", deleted["tenant"] + 1))
|
|
monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: [SimpleNamespace(id="ut-1")])
|
|
monkeypatch.setattr(module.UserTenantService, "delete_by_id", lambda _ut_id: deleted.__setitem__("user_tenant", deleted["user_tenant"] + 1))
|
|
|
|
module.rollback_user_registration("user-1")
|
|
assert deleted == {"user": 1, "tenant": 1, "user_tenant": 1}, deleted
|
|
|
|
monkeypatch.setattr(module.UserService, "delete_by_id", lambda _user_id: (_ for _ in ()).throw(RuntimeError("u boom")))
|
|
monkeypatch.setattr(module.TenantService, "delete_by_id", lambda _tenant_id: (_ for _ in ()).throw(RuntimeError("t boom")))
|
|
monkeypatch.setattr(module.UserTenantService, "query", lambda **_kwargs: (_ for _ in ()).throw(RuntimeError("ut boom")))
|
|
|
|
class _RaisingDeleteQuery:
|
|
def where(self, *_args, **_kwargs):
|
|
raise RuntimeError("llm boom")
|
|
|
|
module.rollback_user_registration("user-2")
|
|
|
|
monkeypatch.setattr(module.UserService, "save", lambda **_kwargs: False)
|
|
res = module.user_register(
|
|
"new-user",
|
|
{
|
|
"nickname": "new",
|
|
"email": "new@example.com",
|
|
"password": "pw",
|
|
"access_token": "tk",
|
|
"login_channel": "password",
|
|
"last_login_time": "2024-01-01 00:00:00",
|
|
"is_superuser": False,
|
|
},
|
|
)
|
|
assert res is None
|
|
|
|
monkeypatch.setattr(module.settings, "REGISTER_ENABLED", False)
|
|
_set_request_json(monkeypatch, module, {"nickname": "neo", "email": "neo@example.com", "password": "enc"})
|
|
res = _run(module.user_add())
|
|
assert res["code"] == module.RetCode.OPERATING_ERROR, res
|
|
assert "disabled" in res["message"], res
|
|
|
|
monkeypatch.setattr(module.settings, "REGISTER_ENABLED", True)
|
|
_set_request_json(monkeypatch, module, {"nickname": "neo", "email": "bad-email", "password": "enc"})
|
|
res = _run(module.user_add())
|
|
assert res["code"] == module.RetCode.OPERATING_ERROR, res
|
|
assert "Invalid email address" in res["message"], res
|
|
|
|
monkeypatch.setattr(module.UserService, "query", lambda **_kwargs: [])
|
|
monkeypatch.setattr(module, "decrypt", lambda value: value)
|
|
monkeypatch.setattr(module, "get_uuid", lambda: "new-user-id")
|
|
rollback_calls = []
|
|
monkeypatch.setattr(module, "rollback_user_registration", lambda user_id: rollback_calls.append(user_id))
|
|
|
|
_set_request_json(monkeypatch, module, {"nickname": "neo", "email": "neo@example.com", "password": "enc"})
|
|
monkeypatch.setattr(module, "user_register", lambda _user_id, _payload: None)
|
|
res = _run(module.user_add())
|
|
assert res["code"] == module.RetCode.EXCEPTION_ERROR, res
|
|
assert "Fail to register neo@example.com." in res["message"], res
|
|
assert rollback_calls == ["new-user-id"], rollback_calls
|
|
|
|
rollback_calls.clear()
|
|
monkeypatch.setattr(
|
|
module,
|
|
"user_register",
|
|
lambda _user_id, _payload: [_DummyUser("dup-1", "neo@example.com"), _DummyUser("dup-2", "neo@example.com")],
|
|
)
|
|
_set_request_json(monkeypatch, module, {"nickname": "neo", "email": "neo@example.com", "password": "enc"})
|
|
res = _run(module.user_add())
|
|
assert res["code"] == module.RetCode.EXCEPTION_ERROR, res
|
|
assert "Same email: neo@example.com exists!" in res["message"], res
|
|
assert rollback_calls == ["new-user-id"], rollback_calls
|
|
|
|
|
|
@pytest.mark.p2
|
|
def test_tenant_info_and_set_tenant_info_exception_matrix_unit(monkeypatch):
|
|
module = _load_user_app(monkeypatch)
|
|
|
|
monkeypatch.setattr(module.TenantService, "get_info_by", lambda _uid: [])
|
|
res = _run(module.tenant_info())
|
|
assert res["code"] == module.RetCode.DATA_ERROR, res
|
|
assert "Tenant not found" in res["message"], res
|
|
|
|
def _raise_tenant_info(_uid):
|
|
raise RuntimeError("tenant info boom")
|
|
|
|
monkeypatch.setattr(module.TenantService, "get_info_by", _raise_tenant_info)
|
|
res = _run(module.tenant_info())
|
|
assert res["code"] == module.RetCode.EXCEPTION_ERROR, res
|
|
assert "tenant info boom" in res["message"], res
|
|
|
|
_set_request_json(
|
|
monkeypatch,
|
|
module,
|
|
{"tenant_id": "tenant-1", "llm_id": "l", "embd_id": "e", "asr_id": "a", "img2txt_id": "i"},
|
|
)
|
|
|
|
def _raise_update(_tenant_id, _payload):
|
|
raise RuntimeError("tenant update boom")
|
|
|
|
monkeypatch.setattr(module.TenantService, "update_by_id", _raise_update)
|
|
res = _run(module.set_tenant_info())
|
|
assert res["code"] == module.RetCode.EXCEPTION_ERROR, res
|
|
assert "tenant update boom" in res["message"], res
|
|
|
|
|
|
@pytest.mark.p2
|
|
def test_forget_captcha_and_send_otp_matrix_unit(monkeypatch):
|
|
module = _load_user_app(monkeypatch)
|
|
|
|
class _Headers(dict):
|
|
def set(self, key, value):
|
|
self[key] = value
|
|
|
|
async def _make_response(data):
|
|
return SimpleNamespace(data=data, headers=_Headers())
|
|
|
|
monkeypatch.setattr(module, "make_response", _make_response)
|
|
|
|
captcha_pkg = ModuleType("captcha")
|
|
captcha_image_mod = ModuleType("captcha.image")
|
|
|
|
class _ImageCaptcha:
|
|
def __init__(self, **_kwargs):
|
|
pass
|
|
|
|
def generate(self, text):
|
|
return SimpleNamespace(read=lambda: f"img:{text}".encode())
|
|
|
|
captcha_image_mod.ImageCaptcha = _ImageCaptcha
|
|
monkeypatch.setitem(sys.modules, "captcha", captcha_pkg)
|
|
monkeypatch.setitem(sys.modules, "captcha.image", captcha_image_mod)
|
|
|
|
_set_request_args(monkeypatch, module, {"email": ""})
|
|
res = _run(module.forget_get_captcha())
|
|
assert res["code"] == module.RetCode.ARGUMENT_ERROR, res
|
|
|
|
monkeypatch.setattr(module.UserService, "query", lambda **_kwargs: [])
|
|
_set_request_args(monkeypatch, module, {"email": "nobody@example.com"})
|
|
res = _run(module.forget_get_captcha())
|
|
assert res["code"] == module.RetCode.DATA_ERROR, res
|
|
|
|
monkeypatch.setattr(module.UserService, "query", lambda **_kwargs: [_DummyUser("u1", "ok@example.com")])
|
|
monkeypatch.setattr(module.secrets, "choice", lambda _allowed: "A")
|
|
_set_request_args(monkeypatch, module, {"email": "ok@example.com"})
|
|
res = _run(module.forget_get_captcha())
|
|
assert res.data.startswith(b"img:"), res
|
|
assert res.headers["Content-Type"] == "image/JPEG", res.headers
|
|
assert module.REDIS_CONN.get(module.captcha_key("ok@example.com")), module.REDIS_CONN.store
|
|
|
|
_set_request_json(monkeypatch, module, {"email": "", "captcha": ""})
|
|
res = _run(module.forget_send_otp())
|
|
assert res["code"] == module.RetCode.ARGUMENT_ERROR, res
|
|
|
|
monkeypatch.setattr(module.UserService, "query", lambda **_kwargs: [])
|
|
_set_request_json(monkeypatch, module, {"email": "none@example.com", "captcha": "AAAA"})
|
|
res = _run(module.forget_send_otp())
|
|
assert res["code"] == module.RetCode.DATA_ERROR, res
|
|
|
|
monkeypatch.setattr(module.UserService, "query", lambda **_kwargs: [_DummyUser("u1", "ok@example.com")])
|
|
_set_request_json(monkeypatch, module, {"email": "ok@example.com", "captcha": "AAAA"})
|
|
module.REDIS_CONN.store.pop(module.captcha_key("ok@example.com"), None)
|
|
res = _run(module.forget_send_otp())
|
|
assert res["code"] == module.RetCode.NOT_EFFECTIVE, res
|
|
|
|
module.REDIS_CONN.store[module.captcha_key("ok@example.com")] = "ABCD"
|
|
_set_request_json(monkeypatch, module, {"email": "ok@example.com", "captcha": "ZZZZ"})
|
|
res = _run(module.forget_send_otp())
|
|
assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res
|
|
|
|
monkeypatch.setattr(module.time, "time", lambda: 1000)
|
|
k_code, k_attempts, k_last, k_lock = module.otp_keys("ok@example.com")
|
|
module.REDIS_CONN.store[module.captcha_key("ok@example.com")] = "ABCD"
|
|
module.REDIS_CONN.store[k_last] = "990"
|
|
_set_request_json(monkeypatch, module, {"email": "ok@example.com", "captcha": "ABCD"})
|
|
res = _run(module.forget_send_otp())
|
|
assert res["code"] == module.RetCode.NOT_EFFECTIVE, res
|
|
assert "wait" in res["message"], res
|
|
|
|
module.REDIS_CONN.store[module.captcha_key("ok@example.com")] = "ABCD"
|
|
module.REDIS_CONN.store[k_last] = "bad-timestamp"
|
|
monkeypatch.setattr(module.secrets, "choice", lambda _allowed: "B")
|
|
monkeypatch.setattr(module.os, "urandom", lambda _n: b"\x00" * 16)
|
|
monkeypatch.setattr(module, "hash_code", lambda code, _salt: f"HASH_{code}")
|
|
|
|
async def _raise_send_email(*_args, **_kwargs):
|
|
raise RuntimeError("send email boom")
|
|
|
|
monkeypatch.setattr(module, "send_email_html", _raise_send_email)
|
|
_set_request_json(monkeypatch, module, {"email": "ok@example.com", "captcha": "ABCD"})
|
|
res = _run(module.forget_send_otp())
|
|
assert res["code"] == module.RetCode.SERVER_ERROR, res
|
|
assert "failed to send email" in res["message"], res
|
|
|
|
async def _ok_send_email(*_args, **_kwargs):
|
|
return True
|
|
|
|
module.REDIS_CONN.store[module.captcha_key("ok@example.com")] = "ABCD"
|
|
module.REDIS_CONN.store.pop(k_last, None)
|
|
monkeypatch.setattr(module, "send_email_html", _ok_send_email)
|
|
_set_request_json(monkeypatch, module, {"email": "ok@example.com", "captcha": "ABCD"})
|
|
res = _run(module.forget_send_otp())
|
|
assert res["code"] == module.RetCode.SUCCESS, res
|
|
assert res["data"] is True, res
|
|
assert module.REDIS_CONN.get(k_code), module.REDIS_CONN.store
|
|
assert module.REDIS_CONN.get(k_attempts) == 0, module.REDIS_CONN.store
|
|
assert module.REDIS_CONN.get(k_lock) is None, module.REDIS_CONN.store
|
|
|
|
|
|
@pytest.mark.p2
|
|
def test_forget_verify_otp_matrix_unit(monkeypatch):
|
|
module = _load_user_app(monkeypatch)
|
|
email = "ok@example.com"
|
|
k_code, k_attempts, k_last, k_lock = module.otp_keys(email)
|
|
salt = b"\x01" * 16
|
|
monkeypatch.setattr(module, "hash_code", lambda code, _salt: f"HASH_{code}")
|
|
|
|
_set_request_json(monkeypatch, module, {})
|
|
res = _run(module.forget_verify_otp())
|
|
assert res["code"] == module.RetCode.ARGUMENT_ERROR, res
|
|
|
|
monkeypatch.setattr(module.UserService, "query", lambda **_kwargs: [])
|
|
_set_request_json(monkeypatch, module, {"email": email, "otp": "ABCDEF"})
|
|
res = _run(module.forget_verify_otp())
|
|
assert res["code"] == module.RetCode.DATA_ERROR, res
|
|
|
|
monkeypatch.setattr(module.UserService, "query", lambda **_kwargs: [_DummyUser("u1", email)])
|
|
module.REDIS_CONN.store[k_lock] = "1"
|
|
_set_request_json(monkeypatch, module, {"email": email, "otp": "ABCDEF"})
|
|
res = _run(module.forget_verify_otp())
|
|
assert res["code"] == module.RetCode.NOT_EFFECTIVE, res
|
|
module.REDIS_CONN.store.pop(k_lock, None)
|
|
|
|
module.REDIS_CONN.store.pop(k_code, None)
|
|
_set_request_json(monkeypatch, module, {"email": email, "otp": "ABCDEF"})
|
|
res = _run(module.forget_verify_otp())
|
|
assert res["code"] == module.RetCode.NOT_EFFECTIVE, res
|
|
|
|
module.REDIS_CONN.store[k_code] = "broken"
|
|
_set_request_json(monkeypatch, module, {"email": email, "otp": "ABCDEF"})
|
|
res = _run(module.forget_verify_otp())
|
|
assert res["code"] == module.RetCode.EXCEPTION_ERROR, res
|
|
|
|
module.REDIS_CONN.store[k_code] = f"HASH_CORRECT:{salt.hex()}"
|
|
module.REDIS_CONN.store[k_attempts] = "bad-int"
|
|
_set_request_json(monkeypatch, module, {"email": email, "otp": "wrong"})
|
|
res = _run(module.forget_verify_otp())
|
|
assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res
|
|
assert module.REDIS_CONN.get(k_attempts) == 1, module.REDIS_CONN.store
|
|
|
|
module.REDIS_CONN.store[k_code] = f"HASH_CORRECT:{salt.hex()}"
|
|
module.REDIS_CONN.store[k_attempts] = str(module.ATTEMPT_LIMIT - 1)
|
|
_set_request_json(monkeypatch, module, {"email": email, "otp": "wrong"})
|
|
res = _run(module.forget_verify_otp())
|
|
assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res
|
|
assert module.REDIS_CONN.get(k_lock) is not None, module.REDIS_CONN.store
|
|
module.REDIS_CONN.store.pop(k_lock, None)
|
|
|
|
module.REDIS_CONN.store[k_code] = f"HASH_ABCDEF:{salt.hex()}"
|
|
module.REDIS_CONN.store[k_attempts] = "0"
|
|
module.REDIS_CONN.store[k_last] = "1000"
|
|
|
|
def _set_with_verified_fail(key, value, _ttl=None):
|
|
if key == module._verified_key(email):
|
|
raise RuntimeError("verified set boom")
|
|
module.REDIS_CONN.store[key] = value
|
|
|
|
monkeypatch.setattr(module.REDIS_CONN, "set", _set_with_verified_fail)
|
|
_set_request_json(monkeypatch, module, {"email": email, "otp": "abcdef"})
|
|
res = _run(module.forget_verify_otp())
|
|
assert res["code"] == module.RetCode.SERVER_ERROR, res
|
|
|
|
monkeypatch.setattr(module.REDIS_CONN, "set", lambda key, value, _ttl=None: module.REDIS_CONN.store.__setitem__(key, value))
|
|
module.REDIS_CONN.store[k_code] = f"HASH_ABCDEF:{salt.hex()}"
|
|
module.REDIS_CONN.store[k_attempts] = "0"
|
|
module.REDIS_CONN.store[k_last] = "1000"
|
|
_set_request_json(monkeypatch, module, {"email": email, "otp": "abcdef"})
|
|
res = _run(module.forget_verify_otp())
|
|
assert res["code"] == module.RetCode.SUCCESS, res
|
|
assert module.REDIS_CONN.get(k_code) is None, module.REDIS_CONN.store
|
|
assert module.REDIS_CONN.get(k_attempts) is None, module.REDIS_CONN.store
|
|
assert module.REDIS_CONN.get(k_last) is None, module.REDIS_CONN.store
|
|
assert module.REDIS_CONN.get(k_lock) is None, module.REDIS_CONN.store
|
|
assert module.REDIS_CONN.get(module._verified_key(email)) == "1", module.REDIS_CONN.store
|
|
|
|
|
|
@pytest.mark.p2
|
|
def test_forget_reset_password_matrix_unit(monkeypatch):
|
|
module = _load_user_app(monkeypatch)
|
|
email = "reset@example.com"
|
|
v_key = module._verified_key(email)
|
|
user = _DummyUser("u-reset", email, nickname="reset-user")
|
|
pwd_a = base64.b64encode(b"new-password").decode()
|
|
pwd_b = base64.b64encode(b"confirm-password").decode()
|
|
pwd_same = base64.b64encode(b"same-password").decode()
|
|
monkeypatch.setattr(module, "decrypt", lambda value: value)
|
|
|
|
_set_request_json(monkeypatch, module, {"email": email, "new_password": pwd_same, "confirm_new_password": pwd_same})
|
|
module.REDIS_CONN.store.pop(v_key, None)
|
|
res = _run(module.forget_reset_password())
|
|
assert res["code"] == module.RetCode.AUTHENTICATION_ERROR, res
|
|
|
|
module.REDIS_CONN.store[v_key] = "1"
|
|
monkeypatch.setattr(module, "decrypt", lambda _value: "")
|
|
_set_request_json(monkeypatch, module, {"email": email, "new_password": "", "confirm_new_password": ""})
|
|
res = _run(module.forget_reset_password())
|
|
assert res["code"] == module.RetCode.ARGUMENT_ERROR, res
|
|
|
|
monkeypatch.setattr(module, "decrypt", lambda value: value)
|
|
module.REDIS_CONN.store[v_key] = "1"
|
|
_set_request_json(monkeypatch, module, {"email": email, "new_password": pwd_a, "confirm_new_password": pwd_b})
|
|
res = _run(module.forget_reset_password())
|
|
assert res["code"] == module.RetCode.ARGUMENT_ERROR, res
|
|
assert "do not match" in res["message"], res
|
|
|
|
module.REDIS_CONN.store[v_key] = "1"
|
|
monkeypatch.setattr(module.UserService, "query_user_by_email", lambda **_kwargs: [])
|
|
_set_request_json(monkeypatch, module, {"email": email, "new_password": pwd_same, "confirm_new_password": pwd_same})
|
|
res = _run(module.forget_reset_password())
|
|
assert res["code"] == module.RetCode.DATA_ERROR, res
|
|
|
|
module.REDIS_CONN.store[v_key] = "1"
|
|
monkeypatch.setattr(module.UserService, "query_user_by_email", lambda **_kwargs: [user])
|
|
|
|
def _raise_update_password(_user_id, _new_pwd):
|
|
raise RuntimeError("reset boom")
|
|
|
|
monkeypatch.setattr(module.UserService, "update_user_password", _raise_update_password)
|
|
_set_request_json(monkeypatch, module, {"email": email, "new_password": pwd_same, "confirm_new_password": pwd_same})
|
|
res = _run(module.forget_reset_password())
|
|
assert res["code"] == module.RetCode.EXCEPTION_ERROR, res
|
|
|
|
module.REDIS_CONN.store[v_key] = "1"
|
|
monkeypatch.setattr(module.UserService, "update_user_password", lambda _user_id, _new_pwd: True)
|
|
monkeypatch.setattr(module.REDIS_CONN, "delete", lambda _key: (_ for _ in ()).throw(RuntimeError("delete boom")))
|
|
_set_request_json(monkeypatch, module, {"email": email, "new_password": pwd_same, "confirm_new_password": pwd_same})
|
|
res = _run(module.forget_reset_password())
|
|
assert res["code"] == module.RetCode.SUCCESS, res
|
|
assert res["auth"] == user.get_id(), res
|
|
|
|
monkeypatch.setattr(module.REDIS_CONN, "delete", lambda key: module.REDIS_CONN.store.pop(key, None))
|
|
module.REDIS_CONN.store[v_key] = "1"
|
|
_set_request_json(monkeypatch, module, {"email": email, "new_password": pwd_same, "confirm_new_password": pwd_same})
|
|
res = _run(module.forget_reset_password())
|
|
assert res["code"] == module.RetCode.SUCCESS, res
|
|
assert res["auth"] == user.get_id(), res
|
|
assert module.REDIS_CONN.get(v_key) is None, module.REDIS_CONN.store
|
|
|
|
|
|
def _load_chat_routes_unit_module(monkeypatch):
|
|
repo_root = Path(__file__).resolve().parents[3]
|
|
module_name = "test_chat_restful_routes_unit_module_for_tenant"
|
|
module_path = repo_root / "api" / "apps" / "restful_apis" / "chat_api.py"
|
|
|
|
quart_mod = ModuleType("quart")
|
|
quart_mod.request = SimpleNamespace(args=SimpleNamespace(get=lambda _key, default=None: default, getlist=lambda _key: []))
|
|
quart_mod.Response = type("_StubResponse", (), {})
|
|
monkeypatch.setitem(sys.modules, "quart", quart_mod)
|
|
|
|
api_pkg = ModuleType("api")
|
|
api_pkg.__path__ = [str(repo_root / "api")]
|
|
monkeypatch.setitem(sys.modules, "api", api_pkg)
|
|
|
|
apps_pkg = ModuleType("api.apps")
|
|
apps_pkg.__path__ = [str(repo_root / "api" / "apps")]
|
|
apps_pkg.current_user = SimpleNamespace(id="tenant-1")
|
|
apps_pkg.login_required = _passthrough_login_required
|
|
monkeypatch.setitem(sys.modules, "api.apps", apps_pkg)
|
|
api_pkg.apps = apps_pkg
|
|
|
|
common_pkg = ModuleType("common")
|
|
common_pkg.__path__ = [str(repo_root / "common")]
|
|
monkeypatch.setitem(sys.modules, "common", common_pkg)
|
|
|
|
settings_mod = ModuleType("common.settings")
|
|
settings_mod.STORAGE_IMPL = type("_StorageImpl", (), {"rm": staticmethod(lambda *_args, **_kwargs: None)})()
|
|
monkeypatch.setitem(sys.modules, "common.settings", settings_mod)
|
|
|
|
constants_mod = ModuleType("common.constants")
|
|
constants_mod.LLMType = SimpleNamespace(CHAT="chat", IMAGE2TEXT="image2text", RERANK="rerank", SPEECH2TEXT="speech2text", TTS="tts")
|
|
constants_mod.RetCode = SimpleNamespace(SUCCESS=0, DATA_ERROR=102, OPERATING_ERROR=103, AUTHENTICATION_ERROR=109)
|
|
constants_mod.StatusEnum = SimpleNamespace(VALID=SimpleNamespace(value="1"), INVALID=SimpleNamespace(value="0"))
|
|
from common.constants import MAXIMUM_PAGE_NUMBER as _MPN, MAXIMUM_TASK_PAGE_NUMBER as _MTPN
|
|
constants_mod.MAXIMUM_PAGE_NUMBER = _MPN
|
|
constants_mod.MAXIMUM_TASK_PAGE_NUMBER = _MTPN
|
|
monkeypatch.setitem(sys.modules, "common.constants", constants_mod)
|
|
|
|
misc_utils_mod = ModuleType("common.misc_utils")
|
|
misc_utils_mod.get_uuid = lambda: "generated-chat-id"
|
|
async def _thread_pool_exec(func, *args, **kwargs):
|
|
return func(*args, **kwargs)
|
|
misc_utils_mod.thread_pool_exec = _thread_pool_exec
|
|
monkeypatch.setitem(sys.modules, "common.misc_utils", misc_utils_mod)
|
|
|
|
dialog_service_mod = ModuleType("api.db.services.dialog_service")
|
|
class _DialogService:
|
|
model = SimpleNamespace(_meta=SimpleNamespace(fields={
|
|
"id": None,
|
|
"tenant_id": None,
|
|
"name": None,
|
|
"description": None,
|
|
"icon": None,
|
|
"kb_ids": None,
|
|
"llm_id": None,
|
|
"llm_setting": None,
|
|
"prompt_config": None,
|
|
"similarity_threshold": None,
|
|
"vector_similarity_weight": None,
|
|
"top_n": None,
|
|
"top_k": None,
|
|
"rerank_id": None,
|
|
"meta_data_filter": None,
|
|
"created_by": None,
|
|
"create_time": None,
|
|
"create_date": None,
|
|
"update_time": None,
|
|
"update_date": None,
|
|
"status": None,
|
|
}))
|
|
@staticmethod
|
|
def query(**_kwargs):
|
|
return []
|
|
@staticmethod
|
|
def save(**_kwargs):
|
|
return True
|
|
@staticmethod
|
|
def get_by_id(_chat_id):
|
|
return False, None
|
|
@staticmethod
|
|
def get_by_tenant_ids(*_args, **_kwargs):
|
|
return [], 0
|
|
dialog_service_mod.DialogService = _DialogService
|
|
dialog_service_mod.async_ask = lambda *_args, **_kwargs: None
|
|
dialog_service_mod.async_chat = lambda *_args, **_kwargs: None
|
|
dialog_service_mod.gen_mindmap = lambda *_args, **_kwargs: None
|
|
monkeypatch.setitem(sys.modules, "api.db.services.dialog_service", dialog_service_mod)
|
|
|
|
conversation_service_mod = ModuleType("api.db.services.conversation_service")
|
|
conversation_service_mod.ConversationService = type("ConversationService", (), {})
|
|
conversation_service_mod.structure_answer = lambda *_args, **_kwargs: {}
|
|
monkeypatch.setitem(sys.modules, "api.db.services.conversation_service", conversation_service_mod)
|
|
|
|
kb_service_mod = ModuleType("api.db.services.knowledgebase_service")
|
|
class _KB:
|
|
def __init__(self):
|
|
self.id = "kb-1"
|
|
self.embd_id = "embd@factory"
|
|
self.chunk_num = 1
|
|
self.name = "Dataset A"
|
|
self.status = "1"
|
|
kb_service_mod.KnowledgebaseService = type('KnowledgebaseService', (), {
|
|
'accessible': staticmethod(lambda **_kwargs: [SimpleNamespace(id='kb-1')]),
|
|
'query': staticmethod(lambda **_kwargs: [_KB()]),
|
|
'get_by_id': staticmethod(lambda _id: (True, _KB())),
|
|
})
|
|
monkeypatch.setitem(sys.modules, "api.db.services.knowledgebase_service", kb_service_mod)
|
|
|
|
tenant_model_provider_mod = ModuleType("api.db.joint_services.tenant_model_service")
|
|
tenant_model_provider_mod.get_model_config_from_provider_instance = lambda *_args, **_kwargs: {}
|
|
tenant_model_provider_mod.get_tenant_default_model_by_type = lambda *_args, **_kwargs: {}
|
|
def _split_model_name(model_name):
|
|
parts = model_name.split("@")
|
|
if len(parts) == 1:
|
|
return parts[0], "", ""
|
|
elif len(parts) == 2:
|
|
return parts[0], "default", parts[1]
|
|
else:
|
|
return parts[0], parts[1], parts[2]
|
|
tenant_model_provider_mod.split_model_name = staticmethod(_split_model_name)
|
|
tenant_model_provider_mod.get_api_key = lambda *_args, **_kwargs: SimpleNamespace(id=1)
|
|
monkeypatch.setitem(sys.modules, "api.db.joint_services.tenant_model_service", tenant_model_provider_mod)
|
|
|
|
llm_service_mod = ModuleType("api.db.services.llm_service")
|
|
llm_service_mod.LLMBundle = lambda *_args, **_kwargs: None
|
|
monkeypatch.setitem(sys.modules, "api.db.services.llm_service", llm_service_mod)
|
|
|
|
search_service_mod = ModuleType("api.db.services.search_service")
|
|
search_service_mod.SearchService = SimpleNamespace()
|
|
monkeypatch.setitem(sys.modules, "api.db.services.search_service", search_service_mod)
|
|
|
|
user_service_mod = ModuleType("api.db.services.user_service")
|
|
user_service_mod.UserService = type('UserService', (), {})
|
|
user_service_mod.TenantService = type('TenantService', (), {
|
|
'get_by_id': staticmethod(lambda _tenant_id: (True, SimpleNamespace(llm_id='glm-4'))),
|
|
'get_joined_tenants_by_user_id': staticmethod(lambda _user_id: [{'tenant_id': 'tenant-1'}, {'tenant_id': 'team-tenant-2'}]),
|
|
})
|
|
user_service_mod.UserTenantService = type('UserTenantService', (), {'query': staticmethod(lambda **_kwargs: [])})
|
|
monkeypatch.setitem(sys.modules, "api.db.services.user_service", user_service_mod)
|
|
|
|
chunk_feedback_service_mod = ModuleType("api.db.services.chunk_feedback_service")
|
|
chunk_feedback_service_mod.ChunkFeedbackService = type('ChunkFeedbackService', (), {'apply_feedback': staticmethod(lambda **_kwargs: {'success_count': 0, 'fail_count': 0, 'chunk_ids': []})})
|
|
monkeypatch.setitem(sys.modules, "api.db.services.chunk_feedback_service", chunk_feedback_service_mod)
|
|
|
|
api_utils_mod = ModuleType("api.utils.api_utils")
|
|
api_utils_mod.check_duplicate_ids = lambda ids, _label: (list(dict.fromkeys(ids or [])), [])
|
|
api_utils_mod.get_data_error_result = lambda message='': {'code': 102, 'data': None, 'message': message}
|
|
api_utils_mod.get_json_result = lambda data=None, message='', code=0: {'code': code, 'data': data, 'message': message}
|
|
api_utils_mod.server_error_response = lambda ex: {'code': 500, 'data': None, 'message': str(ex)}
|
|
api_utils_mod.validate_request = lambda *_args, **_kwargs: (lambda func: func)
|
|
api_utils_mod.get_request_json = lambda: _AwaitableValue({})
|
|
monkeypatch.setitem(sys.modules, "api.utils.api_utils", api_utils_mod)
|
|
|
|
rag_pkg = ModuleType("rag")
|
|
rag_pkg.__path__ = [str(repo_root / 'rag')]
|
|
monkeypatch.setitem(sys.modules, 'rag', rag_pkg)
|
|
rag_prompts_pkg = ModuleType('rag.prompts')
|
|
rag_prompts_pkg.__path__ = [str(repo_root / 'rag' / 'prompts')]
|
|
monkeypatch.setitem(sys.modules, 'rag.prompts', rag_prompts_pkg)
|
|
rag_prompts_generator_mod = ModuleType('rag.prompts.generator')
|
|
rag_prompts_generator_mod.chunks_format = lambda reference: reference.get('chunks', []) if isinstance(reference, dict) else []
|
|
monkeypatch.setitem(sys.modules, 'rag.prompts.generator', rag_prompts_generator_mod)
|
|
rag_prompts_template_mod = ModuleType('rag.prompts.template')
|
|
rag_prompts_template_mod.load_prompt = lambda *_args, **_kwargs: ''
|
|
monkeypatch.setitem(sys.modules, 'rag.prompts.template', rag_prompts_template_mod)
|
|
|
|
spec = importlib.util.spec_from_file_location(module_name, module_path)
|
|
module = importlib.util.module_from_spec(spec)
|
|
module.manager = _DummyManager()
|
|
monkeypatch.setitem(sys.modules, module_name, module)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
@pytest.mark.p1
|
|
def test_create_chat_uses_tenant_default_llm_when_llm_id_is_null_unit(monkeypatch):
|
|
module = _load_chat_routes_unit_module(monkeypatch)
|
|
saved = {}
|
|
|
|
async def _request_json():
|
|
return {
|
|
'name': 'chat-a',
|
|
'dataset_ids': ['kb-1'],
|
|
'llm_id': None,
|
|
'llm_setting': {'temperature': 0.8},
|
|
'prompt_config': {'system': 'Answer with {knowledge}', 'parameters': [{'key': 'knowledge', 'optional': False}]},
|
|
}
|
|
|
|
monkeypatch.setattr(module, 'get_request_json', _request_json)
|
|
monkeypatch.setattr(module.DialogService, 'query', lambda **_kwargs: [])
|
|
|
|
def _save(**kwargs):
|
|
saved.update(kwargs)
|
|
return True
|
|
|
|
monkeypatch.setattr(module.DialogService, 'save', _save)
|
|
monkeypatch.setattr(module.DialogService, 'get_by_id', lambda _id: (True, SimpleNamespace(to_dict=lambda: saved)))
|
|
|
|
res = _run(module.create.__wrapped__())
|
|
assert res['code'] == 0
|
|
assert saved['llm_id'] == 'glm-4'
|
|
assert saved['llm_setting']['temperature'] == 0.8
|
|
|
|
|
|
@pytest.mark.p2
|
|
def test_list_chats_authorized_multi_tenant_unit(monkeypatch):
|
|
module = _load_chat_routes_unit_module(monkeypatch)
|
|
captured = {}
|
|
monkeypatch.setattr(
|
|
module,
|
|
'request',
|
|
SimpleNamespace(
|
|
args=SimpleNamespace(
|
|
get=lambda key, default=None: {
|
|
'keywords': '', 'page': '1', 'page_size': '10', 'orderby': 'create_time', 'desc': 'true', 'id': None, 'name': None,
|
|
}.get(key, default),
|
|
getlist=lambda key: ['tenant-1', 'team-tenant-2'] if key == 'owner_ids' else [],
|
|
)
|
|
),
|
|
)
|
|
|
|
def _get_by_tenant_ids(owner_ids, user_id, *args, **kwargs):
|
|
captured['owner_ids'] = owner_ids
|
|
captured['user_id'] = user_id
|
|
return ([{'id': 'c1', 'tenant_id': 'tenant-1'}, {'id': 'c2', 'tenant_id': 'team-tenant-2'}], 2)
|
|
|
|
monkeypatch.setattr(module.DialogService, 'get_by_tenant_ids', _get_by_tenant_ids)
|
|
monkeypatch.setattr(module.KnowledgebaseService, 'get_by_id', lambda _id: (False, None))
|
|
res = _run(module.list_chats.__wrapped__())
|
|
assert res['code'] == 0
|
|
assert res['data']['total'] == 2
|
|
assert {c['id'] for c in res['data']['chats']} == {'c1', 'c2'}
|
|
assert set(captured['owner_ids']) == {'tenant-1', 'team-tenant-2'}
|
|
assert captured['user_id'] == 'tenant-1'
|