您最多选择25个主题 主题必须以字母或数字开头,可以包含连字符 (-),并且长度不得超过35个字符

test_account_service.py 67KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431143214331434143514361437143814391440144114421443144414451446144714481449145014511452145314541455145614571458145914601461146214631464146514661467146814691470147114721473147414751476147714781479148014811482148314841485148614871488148914901491149214931494149514961497149814991500150115021503150415051506150715081509151015111512151315141515151615171518151915201521152215231524152515261527152815291530153115321533153415351536153715381539154015411542154315441545
  1. import json
  2. from datetime import datetime, timedelta
  3. from unittest.mock import MagicMock, patch
  4. import pytest
  5. from configs import dify_config
  6. from models.account import Account
  7. from services.account_service import AccountService, RegisterService, TenantService
  8. from services.errors.account import (
  9. AccountAlreadyInTenantError,
  10. AccountLoginError,
  11. AccountNotFoundError,
  12. AccountPasswordError,
  13. AccountRegisterError,
  14. CurrentPasswordIncorrectError,
  15. )
  16. from tests.unit_tests.services.services_test_help import ServiceDbTestHelper
  17. class TestAccountAssociatedDataFactory:
  18. """Factory class for creating test data and mock objects for account service tests."""
  19. @staticmethod
  20. def create_account_mock(
  21. account_id: str = "user-123",
  22. email: str = "test@example.com",
  23. name: str = "Test User",
  24. status: str = "active",
  25. password: str = "hashed_password",
  26. password_salt: str = "salt",
  27. interface_language: str = "en-US",
  28. interface_theme: str = "light",
  29. timezone: str = "UTC",
  30. **kwargs,
  31. ) -> MagicMock:
  32. """Create a mock account with specified attributes."""
  33. account = MagicMock(spec=Account)
  34. account.id = account_id
  35. account.email = email
  36. account.name = name
  37. account.status = status
  38. account.password = password
  39. account.password_salt = password_salt
  40. account.interface_language = interface_language
  41. account.interface_theme = interface_theme
  42. account.timezone = timezone
  43. # Set last_active_at to a datetime object that's older than 10 minutes
  44. account.last_active_at = datetime.now() - timedelta(minutes=15)
  45. account.initialized_at = None
  46. for key, value in kwargs.items():
  47. setattr(account, key, value)
  48. return account
  49. @staticmethod
  50. def create_tenant_join_mock(
  51. tenant_id: str = "tenant-456",
  52. account_id: str = "user-123",
  53. current: bool = True,
  54. role: str = "normal",
  55. **kwargs,
  56. ) -> MagicMock:
  57. """Create a mock tenant account join record."""
  58. tenant_join = MagicMock()
  59. tenant_join.tenant_id = tenant_id
  60. tenant_join.account_id = account_id
  61. tenant_join.current = current
  62. tenant_join.role = role
  63. for key, value in kwargs.items():
  64. setattr(tenant_join, key, value)
  65. return tenant_join
  66. @staticmethod
  67. def create_feature_service_mock(allow_register: bool = True):
  68. """Create a mock feature service."""
  69. mock_service = MagicMock()
  70. mock_service.get_system_features.return_value.is_allow_register = allow_register
  71. return mock_service
  72. @staticmethod
  73. def create_billing_service_mock(email_frozen: bool = False):
  74. """Create a mock billing service."""
  75. mock_service = MagicMock()
  76. mock_service.is_email_in_freeze.return_value = email_frozen
  77. return mock_service
  78. class TestAccountService:
  79. """
  80. Comprehensive unit tests for AccountService methods.
  81. This test suite covers all account-related operations including:
  82. - Authentication and login
  83. - Account creation and registration
  84. - Password management
  85. - JWT token generation
  86. - User loading and tenant management
  87. - Error conditions and edge cases
  88. """
  89. @pytest.fixture
  90. def mock_db_dependencies(self):
  91. """Common mock setup for database dependencies."""
  92. with patch("services.account_service.db") as mock_db:
  93. mock_db.session.add = MagicMock()
  94. mock_db.session.commit = MagicMock()
  95. yield {
  96. "db": mock_db,
  97. }
  98. @pytest.fixture
  99. def mock_password_dependencies(self):
  100. """Mock setup for password-related functions."""
  101. with (
  102. patch("services.account_service.compare_password") as mock_compare_password,
  103. patch("services.account_service.hash_password") as mock_hash_password,
  104. patch("services.account_service.valid_password") as mock_valid_password,
  105. ):
  106. yield {
  107. "compare_password": mock_compare_password,
  108. "hash_password": mock_hash_password,
  109. "valid_password": mock_valid_password,
  110. }
  111. @pytest.fixture
  112. def mock_external_service_dependencies(self):
  113. """Mock setup for external service dependencies."""
  114. with (
  115. patch("services.account_service.FeatureService") as mock_feature_service,
  116. patch("services.account_service.BillingService") as mock_billing_service,
  117. patch("services.account_service.PassportService") as mock_passport_service,
  118. ):
  119. yield {
  120. "feature_service": mock_feature_service,
  121. "billing_service": mock_billing_service,
  122. "passport_service": mock_passport_service,
  123. }
  124. @pytest.fixture
  125. def mock_db_with_autospec(self):
  126. """
  127. Mock database with autospec for more realistic behavior.
  128. This approach preserves the actual method signatures and behavior.
  129. """
  130. with patch("services.account_service.db", autospec=True) as mock_db:
  131. # Create a more realistic session mock
  132. mock_session = MagicMock()
  133. mock_db.session = mock_session
  134. # Setup basic session methods
  135. mock_session.add = MagicMock()
  136. mock_session.commit = MagicMock()
  137. mock_session.query = MagicMock()
  138. yield mock_db
  139. def _assert_database_operations_called(self, mock_db):
  140. """Helper method to verify database operations were called."""
  141. mock_db.session.commit.assert_called()
  142. def _assert_database_operations_not_called(self, mock_db):
  143. """Helper method to verify database operations were not called."""
  144. mock_db.session.commit.assert_not_called()
  145. def _assert_exception_raised(self, exception_type, callable_func, *args, **kwargs):
  146. """Helper method to verify that specific exception is raised."""
  147. with pytest.raises(exception_type):
  148. callable_func(*args, **kwargs)
  149. # ==================== Authentication Tests ====================
  150. def test_authenticate_success(self, mock_db_dependencies, mock_password_dependencies):
  151. """Test successful authentication with correct email and password."""
  152. # Setup test data
  153. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  154. # Setup smart database query mock
  155. query_results = {("Account", "email", "test@example.com"): mock_account}
  156. ServiceDbTestHelper.setup_db_query_filter_by_mock(mock_db_dependencies["db"], query_results)
  157. mock_password_dependencies["compare_password"].return_value = True
  158. # Execute test
  159. result = AccountService.authenticate("test@example.com", "password")
  160. # Verify results
  161. assert result == mock_account
  162. self._assert_database_operations_called(mock_db_dependencies["db"])
  163. def test_authenticate_account_not_found(self, mock_db_dependencies):
  164. """Test authentication when account does not exist."""
  165. # Setup smart database query mock - no matching results
  166. query_results = {("Account", "email", "notfound@example.com"): None}
  167. ServiceDbTestHelper.setup_db_query_filter_by_mock(mock_db_dependencies["db"], query_results)
  168. # Execute test and verify exception
  169. self._assert_exception_raised(
  170. AccountNotFoundError, AccountService.authenticate, "notfound@example.com", "password"
  171. )
  172. def test_authenticate_account_banned(self, mock_db_dependencies):
  173. """Test authentication when account is banned."""
  174. # Setup test data
  175. mock_account = TestAccountAssociatedDataFactory.create_account_mock(status="banned")
  176. # Setup smart database query mock
  177. query_results = {("Account", "email", "banned@example.com"): mock_account}
  178. ServiceDbTestHelper.setup_db_query_filter_by_mock(mock_db_dependencies["db"], query_results)
  179. # Execute test and verify exception
  180. self._assert_exception_raised(AccountLoginError, AccountService.authenticate, "banned@example.com", "password")
  181. def test_authenticate_password_error(self, mock_db_dependencies, mock_password_dependencies):
  182. """Test authentication with wrong password."""
  183. # Setup test data
  184. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  185. # Setup smart database query mock
  186. query_results = {("Account", "email", "test@example.com"): mock_account}
  187. ServiceDbTestHelper.setup_db_query_filter_by_mock(mock_db_dependencies["db"], query_results)
  188. mock_password_dependencies["compare_password"].return_value = False
  189. # Execute test and verify exception
  190. self._assert_exception_raised(
  191. AccountPasswordError, AccountService.authenticate, "test@example.com", "wrongpassword"
  192. )
  193. def test_authenticate_pending_account_activates(self, mock_db_dependencies, mock_password_dependencies):
  194. """Test authentication for a pending account, which should activate on login."""
  195. # Setup test data
  196. mock_account = TestAccountAssociatedDataFactory.create_account_mock(status="pending")
  197. # Setup smart database query mock
  198. query_results = {("Account", "email", "pending@example.com"): mock_account}
  199. ServiceDbTestHelper.setup_db_query_filter_by_mock(mock_db_dependencies["db"], query_results)
  200. mock_password_dependencies["compare_password"].return_value = True
  201. # Execute test
  202. result = AccountService.authenticate("pending@example.com", "password")
  203. # Verify results
  204. assert result == mock_account
  205. assert mock_account.status == "active"
  206. self._assert_database_operations_called(mock_db_dependencies["db"])
  207. # ==================== Account Creation Tests ====================
  208. def test_create_account_success(
  209. self, mock_db_dependencies, mock_password_dependencies, mock_external_service_dependencies
  210. ):
  211. """Test successful account creation with all required parameters."""
  212. # Setup mocks
  213. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  214. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  215. mock_password_dependencies["hash_password"].return_value = b"hashed_password"
  216. # Execute test
  217. result = AccountService.create_account(
  218. email="test@example.com",
  219. name="Test User",
  220. interface_language="en-US",
  221. password="password123",
  222. interface_theme="light",
  223. )
  224. # Verify results
  225. assert result.email == "test@example.com"
  226. assert result.name == "Test User"
  227. assert result.interface_language == "en-US"
  228. assert result.interface_theme == "light"
  229. assert result.password is not None
  230. assert result.password_salt is not None
  231. assert result.timezone is not None
  232. # Verify database operations
  233. mock_db_dependencies["db"].session.add.assert_called_once()
  234. added_account = mock_db_dependencies["db"].session.add.call_args[0][0]
  235. assert added_account.email == "test@example.com"
  236. assert added_account.name == "Test User"
  237. assert added_account.interface_language == "en-US"
  238. assert added_account.interface_theme == "light"
  239. assert added_account.password is not None
  240. assert added_account.password_salt is not None
  241. assert added_account.timezone is not None
  242. self._assert_database_operations_called(mock_db_dependencies["db"])
  243. def test_create_account_registration_disabled(self, mock_external_service_dependencies):
  244. """Test account creation when registration is disabled."""
  245. # Setup mocks
  246. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = False
  247. # Execute test and verify exception
  248. self._assert_exception_raised(
  249. Exception, # AccountNotFound
  250. AccountService.create_account,
  251. email="test@example.com",
  252. name="Test User",
  253. interface_language="en-US",
  254. )
  255. def test_create_account_email_frozen(self, mock_db_dependencies, mock_external_service_dependencies):
  256. """Test account creation with frozen email address."""
  257. # Setup mocks
  258. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  259. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = True
  260. dify_config.BILLING_ENABLED = True
  261. # Execute test and verify exception
  262. self._assert_exception_raised(
  263. AccountRegisterError,
  264. AccountService.create_account,
  265. email="frozen@example.com",
  266. name="Test User",
  267. interface_language="en-US",
  268. )
  269. dify_config.BILLING_ENABLED = False
  270. def test_create_account_without_password(self, mock_db_dependencies, mock_external_service_dependencies):
  271. """Test account creation without password (for invite-based registration)."""
  272. # Setup mocks
  273. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  274. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  275. # Execute test
  276. result = AccountService.create_account(
  277. email="test@example.com",
  278. name="Test User",
  279. interface_language="zh-CN",
  280. password=None,
  281. interface_theme="dark",
  282. )
  283. # Verify results
  284. assert result.email == "test@example.com"
  285. assert result.name == "Test User"
  286. assert result.interface_language == "zh-CN"
  287. assert result.interface_theme == "dark"
  288. assert result.password is None
  289. assert result.password_salt is None
  290. assert result.timezone is not None
  291. # Verify database operations
  292. mock_db_dependencies["db"].session.add.assert_called_once()
  293. added_account = mock_db_dependencies["db"].session.add.call_args[0][0]
  294. assert added_account.email == "test@example.com"
  295. assert added_account.name == "Test User"
  296. assert added_account.interface_language == "zh-CN"
  297. assert added_account.interface_theme == "dark"
  298. assert added_account.password is None
  299. assert added_account.password_salt is None
  300. assert added_account.timezone is not None
  301. self._assert_database_operations_called(mock_db_dependencies["db"])
  302. # ==================== Password Management Tests ====================
  303. def test_update_account_password_success(self, mock_db_dependencies, mock_password_dependencies):
  304. """Test successful password update with correct current password and valid new password."""
  305. # Setup test data
  306. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  307. mock_password_dependencies["compare_password"].return_value = True
  308. mock_password_dependencies["valid_password"].return_value = None
  309. mock_password_dependencies["hash_password"].return_value = b"new_hashed_password"
  310. # Execute test
  311. result = AccountService.update_account_password(mock_account, "old_password", "new_password123")
  312. # Verify results
  313. assert result == mock_account
  314. assert mock_account.password is not None
  315. assert mock_account.password_salt is not None
  316. # Verify password validation was called
  317. mock_password_dependencies["compare_password"].assert_called_once_with(
  318. "old_password", "hashed_password", "salt"
  319. )
  320. mock_password_dependencies["valid_password"].assert_called_once_with("new_password123")
  321. # Verify database operations
  322. self._assert_database_operations_called(mock_db_dependencies["db"])
  323. def test_update_account_password_current_password_incorrect(self, mock_password_dependencies):
  324. """Test password update with incorrect current password."""
  325. # Setup test data
  326. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  327. mock_password_dependencies["compare_password"].return_value = False
  328. # Execute test and verify exception
  329. self._assert_exception_raised(
  330. CurrentPasswordIncorrectError,
  331. AccountService.update_account_password,
  332. mock_account,
  333. "wrong_password",
  334. "new_password123",
  335. )
  336. # Verify password comparison was called
  337. mock_password_dependencies["compare_password"].assert_called_once_with(
  338. "wrong_password", "hashed_password", "salt"
  339. )
  340. def test_update_account_password_invalid_new_password(self, mock_password_dependencies):
  341. """Test password update with invalid new password."""
  342. # Setup test data
  343. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  344. mock_password_dependencies["compare_password"].return_value = True
  345. mock_password_dependencies["valid_password"].side_effect = ValueError("Password too short")
  346. # Execute test and verify exception
  347. self._assert_exception_raised(
  348. ValueError, AccountService.update_account_password, mock_account, "old_password", "short"
  349. )
  350. # Verify password validation was called
  351. mock_password_dependencies["valid_password"].assert_called_once_with("short")
  352. # ==================== User Loading Tests ====================
  353. def test_load_user_success(self, mock_db_dependencies):
  354. """Test successful user loading with current tenant."""
  355. # Setup test data
  356. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  357. mock_tenant_join = TestAccountAssociatedDataFactory.create_tenant_join_mock()
  358. # Setup smart database query mock
  359. query_results = {
  360. ("Account", "id", "user-123"): mock_account,
  361. ("TenantAccountJoin", "account_id", "user-123"): mock_tenant_join,
  362. }
  363. ServiceDbTestHelper.setup_db_query_filter_by_mock(mock_db_dependencies["db"], query_results)
  364. # Mock datetime
  365. with patch("services.account_service.datetime") as mock_datetime:
  366. mock_now = datetime.now()
  367. mock_datetime.now.return_value = mock_now
  368. mock_datetime.UTC = "UTC"
  369. # Execute test
  370. result = AccountService.load_user("user-123")
  371. # Verify results
  372. assert result == mock_account
  373. assert mock_account.set_tenant_id.called
  374. def test_load_user_not_found(self, mock_db_dependencies):
  375. """Test user loading when user does not exist."""
  376. # Setup smart database query mock - no matching results
  377. query_results = {("Account", "id", "non-existent-user"): None}
  378. ServiceDbTestHelper.setup_db_query_filter_by_mock(mock_db_dependencies["db"], query_results)
  379. # Execute test
  380. result = AccountService.load_user("non-existent-user")
  381. # Verify results
  382. assert result is None
  383. def test_load_user_banned(self, mock_db_dependencies):
  384. """Test user loading when user is banned."""
  385. # Setup test data
  386. mock_account = TestAccountAssociatedDataFactory.create_account_mock(status="banned")
  387. # Setup smart database query mock
  388. query_results = {("Account", "id", "user-123"): mock_account}
  389. ServiceDbTestHelper.setup_db_query_filter_by_mock(mock_db_dependencies["db"], query_results)
  390. # Execute test and verify exception
  391. self._assert_exception_raised(
  392. Exception, # Unauthorized
  393. AccountService.load_user,
  394. "user-123",
  395. )
  396. def test_load_user_no_current_tenant(self, mock_db_dependencies):
  397. """Test user loading when user has no current tenant but has available tenants."""
  398. # Setup test data
  399. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  400. mock_available_tenant = TestAccountAssociatedDataFactory.create_tenant_join_mock(current=False)
  401. # Setup smart database query mock for complex scenario
  402. query_results = {
  403. ("Account", "id", "user-123"): mock_account,
  404. ("TenantAccountJoin", "account_id", "user-123"): None, # No current tenant
  405. ("TenantAccountJoin", "order_by", "first_available"): mock_available_tenant, # First available tenant
  406. }
  407. ServiceDbTestHelper.setup_db_query_filter_by_mock(mock_db_dependencies["db"], query_results)
  408. # Mock datetime
  409. with patch("services.account_service.datetime") as mock_datetime:
  410. mock_now = datetime.now()
  411. mock_datetime.now.return_value = mock_now
  412. mock_datetime.UTC = "UTC"
  413. # Execute test
  414. result = AccountService.load_user("user-123")
  415. # Verify results
  416. assert result == mock_account
  417. assert mock_available_tenant.current is True
  418. self._assert_database_operations_called(mock_db_dependencies["db"])
  419. def test_load_user_no_tenants(self, mock_db_dependencies):
  420. """Test user loading when user has no tenants at all."""
  421. # Setup test data
  422. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  423. # Setup smart database query mock for no tenants scenario
  424. query_results = {
  425. ("Account", "id", "user-123"): mock_account,
  426. ("TenantAccountJoin", "account_id", "user-123"): None, # No current tenant
  427. ("TenantAccountJoin", "order_by", "first_available"): None, # No available tenants
  428. }
  429. ServiceDbTestHelper.setup_db_query_filter_by_mock(mock_db_dependencies["db"], query_results)
  430. # Mock datetime
  431. with patch("services.account_service.datetime") as mock_datetime:
  432. mock_now = datetime.now()
  433. mock_datetime.now.return_value = mock_now
  434. mock_datetime.UTC = "UTC"
  435. # Execute test
  436. result = AccountService.load_user("user-123")
  437. # Verify results
  438. assert result is None
  439. class TestTenantService:
  440. """
  441. Comprehensive unit tests for TenantService methods.
  442. This test suite covers all tenant-related operations including:
  443. - Tenant creation and management
  444. - Member management and permissions
  445. - Tenant switching
  446. - Role updates and permission checks
  447. - Error conditions and edge cases
  448. """
  449. @pytest.fixture
  450. def mock_db_dependencies(self):
  451. """Common mock setup for database dependencies."""
  452. with patch("services.account_service.db") as mock_db:
  453. mock_db.session.add = MagicMock()
  454. mock_db.session.commit = MagicMock()
  455. yield {
  456. "db": mock_db,
  457. }
  458. @pytest.fixture
  459. def mock_rsa_dependencies(self):
  460. """Mock setup for RSA-related functions."""
  461. with patch("services.account_service.generate_key_pair") as mock_generate_key_pair:
  462. yield mock_generate_key_pair
  463. @pytest.fixture
  464. def mock_external_service_dependencies(self):
  465. """Mock setup for external service dependencies."""
  466. with (
  467. patch("services.account_service.FeatureService") as mock_feature_service,
  468. patch("services.account_service.BillingService") as mock_billing_service,
  469. ):
  470. yield {
  471. "feature_service": mock_feature_service,
  472. "billing_service": mock_billing_service,
  473. }
  474. def _assert_database_operations_called(self, mock_db):
  475. """Helper method to verify database operations were called."""
  476. mock_db.session.commit.assert_called()
  477. def _assert_exception_raised(self, exception_type, callable_func, *args, **kwargs):
  478. """Helper method to verify that specific exception is raised."""
  479. with pytest.raises(exception_type):
  480. callable_func(*args, **kwargs)
  481. # ==================== Tenant Creation Tests ====================
  482. def test_create_owner_tenant_if_not_exist_new_user(
  483. self, mock_db_dependencies, mock_rsa_dependencies, mock_external_service_dependencies
  484. ):
  485. """Test creating owner tenant for new user without existing tenants."""
  486. # Setup test data
  487. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  488. # Setup smart database query mock - no existing tenant joins
  489. query_results = {
  490. ("TenantAccountJoin", "account_id", "user-123"): None,
  491. ("TenantAccountJoin", "tenant_id", "tenant-456"): None, # For has_roles check
  492. }
  493. ServiceDbTestHelper.setup_db_query_filter_by_mock(mock_db_dependencies["db"], query_results)
  494. # Setup external service mocks
  495. mock_external_service_dependencies[
  496. "feature_service"
  497. ].get_system_features.return_value.is_allow_create_workspace = True
  498. mock_external_service_dependencies[
  499. "feature_service"
  500. ].get_system_features.return_value.license.workspaces.is_available.return_value = True
  501. # Mock tenant creation
  502. mock_tenant = MagicMock()
  503. mock_tenant.id = "tenant-456"
  504. mock_tenant.name = "Test User's Workspace"
  505. # Mock database operations
  506. mock_db_dependencies["db"].session.add = MagicMock()
  507. # Mock RSA key generation
  508. mock_rsa_dependencies.return_value = "mock_public_key"
  509. # Mock has_roles method to return False (no existing owner)
  510. with patch("services.account_service.TenantService.has_roles") as mock_has_roles:
  511. mock_has_roles.return_value = False
  512. # Mock Tenant creation to set proper ID
  513. with patch("services.account_service.Tenant") as mock_tenant_class:
  514. mock_tenant_instance = MagicMock()
  515. mock_tenant_instance.id = "tenant-456"
  516. mock_tenant_instance.name = "Test User's Workspace"
  517. mock_tenant_class.return_value = mock_tenant_instance
  518. # Execute test
  519. TenantService.create_owner_tenant_if_not_exist(mock_account)
  520. # Verify tenant was created with correct parameters
  521. mock_db_dependencies["db"].session.add.assert_called()
  522. # Get all calls to session.add
  523. add_calls = mock_db_dependencies["db"].session.add.call_args_list
  524. # Should have at least 2 calls: one for Tenant, one for TenantAccountJoin
  525. assert len(add_calls) >= 2
  526. # Verify Tenant was added with correct name
  527. tenant_added = False
  528. tenant_account_join_added = False
  529. for call in add_calls:
  530. added_object = call[0][0] # First argument of the call
  531. # Check if it's a Tenant object
  532. if hasattr(added_object, "name") and hasattr(added_object, "id"):
  533. # This should be a Tenant object
  534. assert added_object.name == "Test User's Workspace"
  535. tenant_added = True
  536. # Check if it's a TenantAccountJoin object
  537. elif (
  538. hasattr(added_object, "tenant_id")
  539. and hasattr(added_object, "account_id")
  540. and hasattr(added_object, "role")
  541. ):
  542. # This should be a TenantAccountJoin object
  543. assert added_object.tenant_id is not None
  544. assert added_object.account_id == "user-123"
  545. assert added_object.role == "owner"
  546. tenant_account_join_added = True
  547. assert tenant_added, "Tenant object was not added to database"
  548. assert tenant_account_join_added, "TenantAccountJoin object was not added to database"
  549. self._assert_database_operations_called(mock_db_dependencies["db"])
  550. assert mock_rsa_dependencies.called, "RSA key generation was not called"
  551. # ==================== Member Management Tests ====================
  552. def test_create_tenant_member_success(self, mock_db_dependencies):
  553. """Test successful tenant member creation."""
  554. # Setup test data
  555. mock_tenant = MagicMock()
  556. mock_tenant.id = "tenant-456"
  557. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  558. # Setup smart database query mock - no existing member
  559. query_results = {("TenantAccountJoin", "tenant_id", "tenant-456"): None}
  560. ServiceDbTestHelper.setup_db_query_filter_by_mock(mock_db_dependencies["db"], query_results)
  561. # Mock database operations
  562. mock_db_dependencies["db"].session.add = MagicMock()
  563. # Execute test
  564. result = TenantService.create_tenant_member(mock_tenant, mock_account, "normal")
  565. # Verify member was created with correct parameters
  566. assert result is not None
  567. mock_db_dependencies["db"].session.add.assert_called_once()
  568. # Verify the TenantAccountJoin object was added with correct parameters
  569. added_tenant_account_join = mock_db_dependencies["db"].session.add.call_args[0][0]
  570. assert added_tenant_account_join.tenant_id == "tenant-456"
  571. assert added_tenant_account_join.account_id == "user-123"
  572. assert added_tenant_account_join.role == "normal"
  573. self._assert_database_operations_called(mock_db_dependencies["db"])
  574. # ==================== Tenant Switching Tests ====================
  575. def test_switch_tenant_success(self):
  576. """Test successful tenant switching."""
  577. # Setup test data
  578. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  579. mock_tenant_join = TestAccountAssociatedDataFactory.create_tenant_join_mock(
  580. tenant_id="tenant-456", account_id="user-123", current=False
  581. )
  582. # Mock the complex query in switch_tenant method
  583. with patch("services.account_service.db") as mock_db:
  584. # Mock the join query that returns the tenant_account_join
  585. mock_query = MagicMock()
  586. mock_filter = MagicMock()
  587. mock_filter.first.return_value = mock_tenant_join
  588. mock_query.filter.return_value = mock_filter
  589. mock_query.join.return_value = mock_query
  590. mock_db.session.query.return_value = mock_query
  591. # Execute test
  592. TenantService.switch_tenant(mock_account, "tenant-456")
  593. # Verify tenant was switched
  594. assert mock_tenant_join.current is True
  595. self._assert_database_operations_called(mock_db)
  596. def test_switch_tenant_no_tenant_id(self):
  597. """Test tenant switching without providing tenant ID."""
  598. # Setup test data
  599. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  600. # Execute test and verify exception
  601. self._assert_exception_raised(ValueError, TenantService.switch_tenant, mock_account, None)
  602. # ==================== Role Management Tests ====================
  603. def test_update_member_role_success(self):
  604. """Test successful member role update."""
  605. # Setup test data
  606. mock_tenant = MagicMock()
  607. mock_tenant.id = "tenant-456"
  608. mock_member = TestAccountAssociatedDataFactory.create_account_mock(account_id="member-789")
  609. mock_operator = TestAccountAssociatedDataFactory.create_account_mock(account_id="operator-123")
  610. mock_target_join = TestAccountAssociatedDataFactory.create_tenant_join_mock(
  611. tenant_id="tenant-456", account_id="member-789", role="normal"
  612. )
  613. mock_operator_join = TestAccountAssociatedDataFactory.create_tenant_join_mock(
  614. tenant_id="tenant-456", account_id="operator-123", role="owner"
  615. )
  616. # Mock the database queries in update_member_role method
  617. with patch("services.account_service.db") as mock_db:
  618. # Mock the first query for operator permission check
  619. mock_query1 = MagicMock()
  620. mock_filter1 = MagicMock()
  621. mock_filter1.first.return_value = mock_operator_join
  622. mock_query1.filter_by.return_value = mock_filter1
  623. # Mock the second query for target member
  624. mock_query2 = MagicMock()
  625. mock_filter2 = MagicMock()
  626. mock_filter2.first.return_value = mock_target_join
  627. mock_query2.filter_by.return_value = mock_filter2
  628. # Make the query method return different mocks for different calls
  629. mock_db.session.query.side_effect = [mock_query1, mock_query2]
  630. # Execute test
  631. TenantService.update_member_role(mock_tenant, mock_member, "admin", mock_operator)
  632. # Verify role was updated
  633. assert mock_target_join.role == "admin"
  634. self._assert_database_operations_called(mock_db)
  635. # ==================== Permission Check Tests ====================
  636. def test_check_member_permission_success(self, mock_db_dependencies):
  637. """Test successful member permission check."""
  638. # Setup test data
  639. mock_tenant = MagicMock()
  640. mock_tenant.id = "tenant-456"
  641. mock_operator = TestAccountAssociatedDataFactory.create_account_mock(account_id="operator-123")
  642. mock_member = TestAccountAssociatedDataFactory.create_account_mock(account_id="member-789")
  643. mock_operator_join = TestAccountAssociatedDataFactory.create_tenant_join_mock(
  644. tenant_id="tenant-456", account_id="operator-123", role="owner"
  645. )
  646. # Setup smart database query mock
  647. query_results = {("TenantAccountJoin", "tenant_id", "tenant-456"): mock_operator_join}
  648. ServiceDbTestHelper.setup_db_query_filter_by_mock(mock_db_dependencies["db"], query_results)
  649. # Execute test - should not raise exception
  650. TenantService.check_member_permission(mock_tenant, mock_operator, mock_member, "add")
  651. def test_check_member_permission_operate_self(self):
  652. """Test member permission check when operator tries to operate self."""
  653. # Setup test data
  654. mock_tenant = MagicMock()
  655. mock_tenant.id = "tenant-456"
  656. mock_operator = TestAccountAssociatedDataFactory.create_account_mock(account_id="operator-123")
  657. # Execute test and verify exception
  658. from services.errors.account import CannotOperateSelfError
  659. self._assert_exception_raised(
  660. CannotOperateSelfError,
  661. TenantService.check_member_permission,
  662. mock_tenant,
  663. mock_operator,
  664. mock_operator, # Same as operator
  665. "add",
  666. )
  667. class TestRegisterService:
  668. """
  669. Comprehensive unit tests for RegisterService methods.
  670. This test suite covers all registration-related operations including:
  671. - System setup
  672. - Account registration
  673. - Member invitation
  674. - Token management
  675. - Invitation validation
  676. - Error conditions and edge cases
  677. """
  678. @pytest.fixture
  679. def mock_db_dependencies(self):
  680. """Common mock setup for database dependencies."""
  681. with patch("services.account_service.db") as mock_db:
  682. mock_db.session.add = MagicMock()
  683. mock_db.session.commit = MagicMock()
  684. mock_db.session.begin_nested = MagicMock()
  685. mock_db.session.rollback = MagicMock()
  686. yield {
  687. "db": mock_db,
  688. }
  689. @pytest.fixture
  690. def mock_redis_dependencies(self):
  691. """Mock setup for Redis-related functions."""
  692. with patch("services.account_service.redis_client") as mock_redis:
  693. yield mock_redis
  694. @pytest.fixture
  695. def mock_external_service_dependencies(self):
  696. """Mock setup for external service dependencies."""
  697. with (
  698. patch("services.account_service.FeatureService") as mock_feature_service,
  699. patch("services.account_service.BillingService") as mock_billing_service,
  700. patch("services.account_service.PassportService") as mock_passport_service,
  701. ):
  702. yield {
  703. "feature_service": mock_feature_service,
  704. "billing_service": mock_billing_service,
  705. "passport_service": mock_passport_service,
  706. }
  707. @pytest.fixture
  708. def mock_task_dependencies(self):
  709. """Mock setup for task dependencies."""
  710. with patch("services.account_service.send_invite_member_mail_task") as mock_send_mail:
  711. yield mock_send_mail
  712. def _assert_database_operations_called(self, mock_db):
  713. """Helper method to verify database operations were called."""
  714. mock_db.session.commit.assert_called()
  715. def _assert_database_operations_not_called(self, mock_db):
  716. """Helper method to verify database operations were not called."""
  717. mock_db.session.commit.assert_not_called()
  718. def _assert_exception_raised(self, exception_type, callable_func, *args, **kwargs):
  719. """Helper method to verify that specific exception is raised."""
  720. with pytest.raises(exception_type):
  721. callable_func(*args, **kwargs)
  722. # ==================== Setup Tests ====================
  723. def test_setup_success(self, mock_db_dependencies, mock_external_service_dependencies):
  724. """Test successful system setup."""
  725. # Setup mocks
  726. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  727. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  728. # Mock AccountService.create_account
  729. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  730. with patch("services.account_service.AccountService.create_account") as mock_create_account:
  731. mock_create_account.return_value = mock_account
  732. # Mock TenantService.create_owner_tenant_if_not_exist
  733. with patch("services.account_service.TenantService.create_owner_tenant_if_not_exist") as mock_create_tenant:
  734. # Mock DifySetup
  735. with patch("services.account_service.DifySetup") as mock_dify_setup:
  736. mock_dify_setup_instance = MagicMock()
  737. mock_dify_setup.return_value = mock_dify_setup_instance
  738. # Execute test
  739. RegisterService.setup("admin@example.com", "Admin User", "password123", "192.168.1.1")
  740. # Verify results
  741. mock_create_account.assert_called_once_with(
  742. email="admin@example.com",
  743. name="Admin User",
  744. interface_language="en-US",
  745. password="password123",
  746. is_setup=True,
  747. )
  748. mock_create_tenant.assert_called_once_with(account=mock_account, is_setup=True)
  749. mock_dify_setup.assert_called_once()
  750. self._assert_database_operations_called(mock_db_dependencies["db"])
  751. def test_setup_failure_rollback(self, mock_db_dependencies, mock_external_service_dependencies):
  752. """Test setup failure with proper rollback."""
  753. # Setup mocks to simulate failure
  754. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  755. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  756. # Mock AccountService.create_account to raise exception
  757. with patch("services.account_service.AccountService.create_account") as mock_create_account:
  758. mock_create_account.side_effect = Exception("Database error")
  759. # Execute test and verify exception
  760. self._assert_exception_raised(
  761. ValueError,
  762. RegisterService.setup,
  763. "admin@example.com",
  764. "Admin User",
  765. "password123",
  766. "192.168.1.1",
  767. )
  768. # Verify rollback operations were called
  769. mock_db_dependencies["db"].session.query.assert_called()
  770. # ==================== Registration Tests ====================
  771. def test_register_success(self, mock_db_dependencies, mock_external_service_dependencies):
  772. """Test successful account registration."""
  773. # Setup mocks
  774. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  775. mock_external_service_dependencies[
  776. "feature_service"
  777. ].get_system_features.return_value.is_allow_create_workspace = True
  778. mock_external_service_dependencies[
  779. "feature_service"
  780. ].get_system_features.return_value.license.workspaces.is_available.return_value = True
  781. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  782. # Mock AccountService.create_account
  783. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  784. with patch("services.account_service.AccountService.create_account") as mock_create_account:
  785. mock_create_account.return_value = mock_account
  786. # Mock TenantService.create_tenant and create_tenant_member
  787. with (
  788. patch("services.account_service.TenantService.create_tenant") as mock_create_tenant,
  789. patch("services.account_service.TenantService.create_tenant_member") as mock_create_member,
  790. patch("services.account_service.tenant_was_created") as mock_event,
  791. ):
  792. mock_tenant = MagicMock()
  793. mock_tenant.id = "tenant-456"
  794. mock_create_tenant.return_value = mock_tenant
  795. # Execute test
  796. result = RegisterService.register(
  797. email="test@example.com",
  798. name="Test User",
  799. password="password123",
  800. language="en-US",
  801. )
  802. # Verify results
  803. assert result == mock_account
  804. assert result.status == "active"
  805. assert result.initialized_at is not None
  806. mock_create_account.assert_called_once_with(
  807. email="test@example.com",
  808. name="Test User",
  809. interface_language="en-US",
  810. password="password123",
  811. is_setup=False,
  812. )
  813. mock_create_tenant.assert_called_once_with("Test User's Workspace")
  814. mock_create_member.assert_called_once_with(mock_tenant, mock_account, role="owner")
  815. mock_event.send.assert_called_once_with(mock_tenant)
  816. self._assert_database_operations_called(mock_db_dependencies["db"])
  817. def test_register_with_oauth(self, mock_db_dependencies, mock_external_service_dependencies):
  818. """Test account registration with OAuth integration."""
  819. # Setup mocks
  820. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  821. mock_external_service_dependencies[
  822. "feature_service"
  823. ].get_system_features.return_value.is_allow_create_workspace = True
  824. mock_external_service_dependencies[
  825. "feature_service"
  826. ].get_system_features.return_value.license.workspaces.is_available.return_value = True
  827. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  828. # Mock AccountService.create_account and link_account_integrate
  829. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  830. with (
  831. patch("services.account_service.AccountService.create_account") as mock_create_account,
  832. patch("services.account_service.AccountService.link_account_integrate") as mock_link_account,
  833. ):
  834. mock_create_account.return_value = mock_account
  835. # Mock TenantService methods
  836. with (
  837. patch("services.account_service.TenantService.create_tenant") as mock_create_tenant,
  838. patch("services.account_service.TenantService.create_tenant_member") as mock_create_member,
  839. patch("services.account_service.tenant_was_created") as mock_event,
  840. ):
  841. mock_tenant = MagicMock()
  842. mock_create_tenant.return_value = mock_tenant
  843. # Execute test
  844. result = RegisterService.register(
  845. email="test@example.com",
  846. name="Test User",
  847. password=None,
  848. open_id="oauth123",
  849. provider="google",
  850. language="en-US",
  851. )
  852. # Verify results
  853. assert result == mock_account
  854. mock_link_account.assert_called_once_with("google", "oauth123", mock_account)
  855. self._assert_database_operations_called(mock_db_dependencies["db"])
  856. def test_register_with_pending_status(self, mock_db_dependencies, mock_external_service_dependencies):
  857. """Test account registration with pending status."""
  858. # Setup mocks
  859. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  860. mock_external_service_dependencies[
  861. "feature_service"
  862. ].get_system_features.return_value.is_allow_create_workspace = True
  863. mock_external_service_dependencies[
  864. "feature_service"
  865. ].get_system_features.return_value.license.workspaces.is_available.return_value = True
  866. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  867. # Mock AccountService.create_account
  868. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  869. with patch("services.account_service.AccountService.create_account") as mock_create_account:
  870. mock_create_account.return_value = mock_account
  871. # Mock TenantService methods
  872. with (
  873. patch("services.account_service.TenantService.create_tenant") as mock_create_tenant,
  874. patch("services.account_service.TenantService.create_tenant_member") as mock_create_member,
  875. patch("services.account_service.tenant_was_created") as mock_event,
  876. ):
  877. mock_tenant = MagicMock()
  878. mock_create_tenant.return_value = mock_tenant
  879. # Execute test with pending status
  880. from models.account import AccountStatus
  881. result = RegisterService.register(
  882. email="test@example.com",
  883. name="Test User",
  884. password="password123",
  885. language="en-US",
  886. status=AccountStatus.PENDING,
  887. )
  888. # Verify results
  889. assert result == mock_account
  890. assert result.status == "pending"
  891. self._assert_database_operations_called(mock_db_dependencies["db"])
  892. def test_register_workspace_not_allowed(self, mock_db_dependencies, mock_external_service_dependencies):
  893. """Test registration when workspace creation is not allowed."""
  894. # Setup mocks
  895. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  896. mock_external_service_dependencies[
  897. "feature_service"
  898. ].get_system_features.return_value.is_allow_create_workspace = True
  899. mock_external_service_dependencies[
  900. "feature_service"
  901. ].get_system_features.return_value.license.workspaces.is_available.return_value = True
  902. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  903. # Mock AccountService.create_account
  904. mock_account = TestAccountAssociatedDataFactory.create_account_mock()
  905. with patch("services.account_service.AccountService.create_account") as mock_create_account:
  906. mock_create_account.return_value = mock_account
  907. # Execute test and verify exception
  908. from services.errors.workspace import WorkSpaceNotAllowedCreateError
  909. with patch("services.account_service.TenantService.create_tenant") as mock_create_tenant:
  910. mock_create_tenant.side_effect = WorkSpaceNotAllowedCreateError()
  911. self._assert_exception_raised(
  912. AccountRegisterError,
  913. RegisterService.register,
  914. email="test@example.com",
  915. name="Test User",
  916. password="password123",
  917. language="en-US",
  918. )
  919. # Verify rollback was called
  920. mock_db_dependencies["db"].session.rollback.assert_called()
  921. def test_register_general_exception(self, mock_db_dependencies, mock_external_service_dependencies):
  922. """Test registration with general exception handling."""
  923. # Setup mocks
  924. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  925. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  926. # Mock AccountService.create_account to raise exception
  927. with patch("services.account_service.AccountService.create_account") as mock_create_account:
  928. mock_create_account.side_effect = Exception("Unexpected error")
  929. # Execute test and verify exception
  930. self._assert_exception_raised(
  931. AccountRegisterError,
  932. RegisterService.register,
  933. email="test@example.com",
  934. name="Test User",
  935. password="password123",
  936. language="en-US",
  937. )
  938. # Verify rollback was called
  939. mock_db_dependencies["db"].session.rollback.assert_called()
  940. # ==================== Member Invitation Tests ====================
  941. def test_invite_new_member_new_account(self, mock_db_dependencies, mock_redis_dependencies, mock_task_dependencies):
  942. """Test inviting a new member who doesn't have an account."""
  943. # Setup test data
  944. mock_tenant = MagicMock()
  945. mock_tenant.id = "tenant-456"
  946. mock_tenant.name = "Test Workspace"
  947. mock_inviter = TestAccountAssociatedDataFactory.create_account_mock(account_id="inviter-123", name="Inviter")
  948. # Mock database queries - need to mock the Session query
  949. mock_session = MagicMock()
  950. mock_session.query.return_value.filter_by.return_value.first.return_value = None # No existing account
  951. with patch("services.account_service.Session") as mock_session_class:
  952. mock_session_class.return_value.__enter__.return_value = mock_session
  953. mock_session_class.return_value.__exit__.return_value = None
  954. # Mock RegisterService.register
  955. mock_new_account = TestAccountAssociatedDataFactory.create_account_mock(
  956. account_id="new-user-456", email="newuser@example.com", name="newuser", status="pending"
  957. )
  958. with patch("services.account_service.RegisterService.register") as mock_register:
  959. mock_register.return_value = mock_new_account
  960. # Mock TenantService methods
  961. with (
  962. patch("services.account_service.TenantService.check_member_permission") as mock_check_permission,
  963. patch("services.account_service.TenantService.create_tenant_member") as mock_create_member,
  964. patch("services.account_service.TenantService.switch_tenant") as mock_switch_tenant,
  965. patch("services.account_service.RegisterService.generate_invite_token") as mock_generate_token,
  966. ):
  967. mock_generate_token.return_value = "invite-token-123"
  968. # Execute test
  969. result = RegisterService.invite_new_member(
  970. tenant=mock_tenant,
  971. email="newuser@example.com",
  972. language="en-US",
  973. role="normal",
  974. inviter=mock_inviter,
  975. )
  976. # Verify results
  977. assert result == "invite-token-123"
  978. mock_register.assert_called_once_with(
  979. email="newuser@example.com",
  980. name="newuser",
  981. language="en-US",
  982. status="pending",
  983. is_setup=True,
  984. )
  985. mock_create_member.assert_called_once_with(mock_tenant, mock_new_account, "normal")
  986. mock_switch_tenant.assert_called_once_with(mock_new_account, mock_tenant.id)
  987. mock_generate_token.assert_called_once_with(mock_tenant, mock_new_account)
  988. mock_task_dependencies.delay.assert_called_once()
  989. def test_invite_new_member_existing_account(
  990. self, mock_db_dependencies, mock_redis_dependencies, mock_task_dependencies
  991. ):
  992. """Test inviting a new member who already has an account."""
  993. # Setup test data
  994. mock_tenant = MagicMock()
  995. mock_tenant.id = "tenant-456"
  996. mock_tenant.name = "Test Workspace"
  997. mock_inviter = TestAccountAssociatedDataFactory.create_account_mock(account_id="inviter-123", name="Inviter")
  998. mock_existing_account = TestAccountAssociatedDataFactory.create_account_mock(
  999. account_id="existing-user-456", email="existing@example.com", status="pending"
  1000. )
  1001. # Mock database queries - need to mock the Session query
  1002. mock_session = MagicMock()
  1003. mock_session.query.return_value.filter_by.return_value.first.return_value = mock_existing_account
  1004. with patch("services.account_service.Session") as mock_session_class:
  1005. mock_session_class.return_value.__enter__.return_value = mock_session
  1006. mock_session_class.return_value.__exit__.return_value = None
  1007. # Mock the db.session.query for TenantAccountJoin
  1008. mock_db_query = MagicMock()
  1009. mock_db_query.filter_by.return_value.first.return_value = None # No existing member
  1010. mock_db_dependencies["db"].session.query.return_value = mock_db_query
  1011. # Mock TenantService methods
  1012. with (
  1013. patch("services.account_service.TenantService.check_member_permission") as mock_check_permission,
  1014. patch("services.account_service.TenantService.create_tenant_member") as mock_create_member,
  1015. patch("services.account_service.RegisterService.generate_invite_token") as mock_generate_token,
  1016. ):
  1017. mock_generate_token.return_value = "invite-token-123"
  1018. # Execute test
  1019. result = RegisterService.invite_new_member(
  1020. tenant=mock_tenant,
  1021. email="existing@example.com",
  1022. language="en-US",
  1023. role="normal",
  1024. inviter=mock_inviter,
  1025. )
  1026. # Verify results
  1027. assert result == "invite-token-123"
  1028. mock_create_member.assert_called_once_with(mock_tenant, mock_existing_account, "normal")
  1029. mock_generate_token.assert_called_once_with(mock_tenant, mock_existing_account)
  1030. mock_task_dependencies.delay.assert_called_once()
  1031. def test_invite_new_member_already_in_tenant(self, mock_db_dependencies, mock_redis_dependencies):
  1032. """Test inviting a member who is already in the tenant."""
  1033. # Setup test data
  1034. mock_tenant = MagicMock()
  1035. mock_tenant.id = "tenant-456"
  1036. mock_inviter = TestAccountAssociatedDataFactory.create_account_mock(account_id="inviter-123", name="Inviter")
  1037. mock_existing_account = TestAccountAssociatedDataFactory.create_account_mock(
  1038. account_id="existing-user-456", email="existing@example.com", status="active"
  1039. )
  1040. # Mock database queries
  1041. query_results = {
  1042. ("Account", "email", "existing@example.com"): mock_existing_account,
  1043. (
  1044. "TenantAccountJoin",
  1045. "tenant_id",
  1046. "tenant-456",
  1047. ): TestAccountAssociatedDataFactory.create_tenant_join_mock(),
  1048. }
  1049. ServiceDbTestHelper.setup_db_query_filter_by_mock(mock_db_dependencies["db"], query_results)
  1050. # Mock TenantService methods
  1051. with patch("services.account_service.TenantService.check_member_permission") as mock_check_permission:
  1052. # Execute test and verify exception
  1053. self._assert_exception_raised(
  1054. AccountAlreadyInTenantError,
  1055. RegisterService.invite_new_member,
  1056. tenant=mock_tenant,
  1057. email="existing@example.com",
  1058. language="en-US",
  1059. role="normal",
  1060. inviter=mock_inviter,
  1061. )
  1062. def test_invite_new_member_no_inviter(self):
  1063. """Test inviting a member without providing an inviter."""
  1064. # Setup test data
  1065. mock_tenant = MagicMock()
  1066. # Execute test and verify exception
  1067. self._assert_exception_raised(
  1068. ValueError,
  1069. RegisterService.invite_new_member,
  1070. tenant=mock_tenant,
  1071. email="test@example.com",
  1072. language="en-US",
  1073. role="normal",
  1074. inviter=None,
  1075. )
  1076. # ==================== Token Management Tests ====================
  1077. def test_generate_invite_token_success(self, mock_redis_dependencies):
  1078. """Test successful invite token generation."""
  1079. # Setup test data
  1080. mock_tenant = MagicMock()
  1081. mock_tenant.id = "tenant-456"
  1082. mock_account = TestAccountAssociatedDataFactory.create_account_mock(
  1083. account_id="user-123", email="test@example.com"
  1084. )
  1085. # Mock uuid generation
  1086. with patch("services.account_service.uuid.uuid4") as mock_uuid:
  1087. mock_uuid.return_value = "test-uuid-123"
  1088. # Execute test
  1089. result = RegisterService.generate_invite_token(mock_tenant, mock_account)
  1090. # Verify results
  1091. assert result == "test-uuid-123"
  1092. mock_redis_dependencies.setex.assert_called_once()
  1093. # Verify the stored data
  1094. call_args = mock_redis_dependencies.setex.call_args
  1095. assert call_args[0][0] == "member_invite:token:test-uuid-123"
  1096. stored_data = json.loads(call_args[0][2])
  1097. assert stored_data["account_id"] == "user-123"
  1098. assert stored_data["email"] == "test@example.com"
  1099. assert stored_data["workspace_id"] == "tenant-456"
  1100. def test_is_valid_invite_token_valid(self, mock_redis_dependencies):
  1101. """Test checking valid invite token."""
  1102. # Setup mock
  1103. mock_redis_dependencies.get.return_value = b'{"test": "data"}'
  1104. # Execute test
  1105. result = RegisterService.is_valid_invite_token("valid-token")
  1106. # Verify results
  1107. assert result is True
  1108. mock_redis_dependencies.get.assert_called_once_with("member_invite:token:valid-token")
  1109. def test_is_valid_invite_token_invalid(self, mock_redis_dependencies):
  1110. """Test checking invalid invite token."""
  1111. # Setup mock
  1112. mock_redis_dependencies.get.return_value = None
  1113. # Execute test
  1114. result = RegisterService.is_valid_invite_token("invalid-token")
  1115. # Verify results
  1116. assert result is False
  1117. mock_redis_dependencies.get.assert_called_once_with("member_invite:token:invalid-token")
  1118. def test_revoke_token_with_workspace_and_email(self, mock_redis_dependencies):
  1119. """Test revoking token with workspace ID and email."""
  1120. # Execute test
  1121. RegisterService.revoke_token("workspace-123", "test@example.com", "token-123")
  1122. # Verify results
  1123. mock_redis_dependencies.delete.assert_called_once()
  1124. call_args = mock_redis_dependencies.delete.call_args
  1125. assert "workspace-123" in call_args[0][0]
  1126. # The email is hashed, so we check for the hash pattern instead
  1127. assert "member_invite_token:" in call_args[0][0]
  1128. def test_revoke_token_without_workspace_and_email(self, mock_redis_dependencies):
  1129. """Test revoking token without workspace ID and email."""
  1130. # Execute test
  1131. RegisterService.revoke_token("", "", "token-123")
  1132. # Verify results
  1133. mock_redis_dependencies.delete.assert_called_once_with("member_invite:token:token-123")
  1134. # ==================== Invitation Validation Tests ====================
  1135. def test_get_invitation_if_token_valid_success(self, mock_db_dependencies, mock_redis_dependencies):
  1136. """Test successful invitation validation."""
  1137. # Setup test data
  1138. mock_tenant = MagicMock()
  1139. mock_tenant.id = "tenant-456"
  1140. mock_tenant.status = "normal"
  1141. mock_account = TestAccountAssociatedDataFactory.create_account_mock(
  1142. account_id="user-123", email="test@example.com"
  1143. )
  1144. with patch("services.account_service.RegisterService._get_invitation_by_token") as mock_get_invitation_by_token:
  1145. # Mock the invitation data returned by _get_invitation_by_token
  1146. invitation_data = {
  1147. "account_id": "user-123",
  1148. "email": "test@example.com",
  1149. "workspace_id": "tenant-456",
  1150. }
  1151. mock_get_invitation_by_token.return_value = invitation_data
  1152. # Mock database queries - complex query mocking
  1153. mock_query1 = MagicMock()
  1154. mock_query1.filter.return_value.first.return_value = mock_tenant
  1155. mock_query2 = MagicMock()
  1156. mock_query2.join.return_value.filter.return_value.first.return_value = (mock_account, "normal")
  1157. mock_db_dependencies["db"].session.query.side_effect = [mock_query1, mock_query2]
  1158. # Execute test
  1159. result = RegisterService.get_invitation_if_token_valid("tenant-456", "test@example.com", "token-123")
  1160. # Verify results
  1161. assert result is not None
  1162. assert result["account"] == mock_account
  1163. assert result["tenant"] == mock_tenant
  1164. assert result["data"] == invitation_data
  1165. def test_get_invitation_if_token_valid_no_token_data(self, mock_redis_dependencies):
  1166. """Test invitation validation with no token data."""
  1167. # Setup mock
  1168. mock_redis_dependencies.get.return_value = None
  1169. # Execute test
  1170. result = RegisterService.get_invitation_if_token_valid("tenant-456", "test@example.com", "token-123")
  1171. # Verify results
  1172. assert result is None
  1173. def test_get_invitation_if_token_valid_tenant_not_found(self, mock_db_dependencies, mock_redis_dependencies):
  1174. """Test invitation validation when tenant is not found."""
  1175. # Setup mock Redis data
  1176. invitation_data = {
  1177. "account_id": "user-123",
  1178. "email": "test@example.com",
  1179. "workspace_id": "tenant-456",
  1180. }
  1181. mock_redis_dependencies.get.return_value = json.dumps(invitation_data).encode()
  1182. # Mock database queries - no tenant found
  1183. mock_query = MagicMock()
  1184. mock_query.filter.return_value.first.return_value = None
  1185. mock_db_dependencies["db"].session.query.return_value = mock_query
  1186. # Execute test
  1187. result = RegisterService.get_invitation_if_token_valid("tenant-456", "test@example.com", "token-123")
  1188. # Verify results
  1189. assert result is None
  1190. def test_get_invitation_if_token_valid_account_not_found(self, mock_db_dependencies, mock_redis_dependencies):
  1191. """Test invitation validation when account is not found."""
  1192. # Setup test data
  1193. mock_tenant = MagicMock()
  1194. mock_tenant.id = "tenant-456"
  1195. mock_tenant.status = "normal"
  1196. # Mock Redis data
  1197. invitation_data = {
  1198. "account_id": "user-123",
  1199. "email": "test@example.com",
  1200. "workspace_id": "tenant-456",
  1201. }
  1202. mock_redis_dependencies.get.return_value = json.dumps(invitation_data).encode()
  1203. # Mock database queries
  1204. mock_query1 = MagicMock()
  1205. mock_query1.filter.return_value.first.return_value = mock_tenant
  1206. mock_query2 = MagicMock()
  1207. mock_query2.join.return_value.filter.return_value.first.return_value = None # No account found
  1208. mock_db_dependencies["db"].session.query.side_effect = [mock_query1, mock_query2]
  1209. # Execute test
  1210. result = RegisterService.get_invitation_if_token_valid("tenant-456", "test@example.com", "token-123")
  1211. # Verify results
  1212. assert result is None
  1213. def test_get_invitation_if_token_valid_account_id_mismatch(self, mock_db_dependencies, mock_redis_dependencies):
  1214. """Test invitation validation when account ID doesn't match."""
  1215. # Setup test data
  1216. mock_tenant = MagicMock()
  1217. mock_tenant.id = "tenant-456"
  1218. mock_tenant.status = "normal"
  1219. mock_account = TestAccountAssociatedDataFactory.create_account_mock(
  1220. account_id="different-user-456", email="test@example.com"
  1221. )
  1222. # Mock Redis data with different account ID
  1223. invitation_data = {
  1224. "account_id": "user-123",
  1225. "email": "test@example.com",
  1226. "workspace_id": "tenant-456",
  1227. }
  1228. mock_redis_dependencies.get.return_value = json.dumps(invitation_data).encode()
  1229. # Mock database queries
  1230. mock_query1 = MagicMock()
  1231. mock_query1.filter.return_value.first.return_value = mock_tenant
  1232. mock_query2 = MagicMock()
  1233. mock_query2.join.return_value.filter.return_value.first.return_value = (mock_account, "normal")
  1234. mock_db_dependencies["db"].session.query.side_effect = [mock_query1, mock_query2]
  1235. # Execute test
  1236. result = RegisterService.get_invitation_if_token_valid("tenant-456", "test@example.com", "token-123")
  1237. # Verify results
  1238. assert result is None
  1239. # ==================== Helper Method Tests ====================
  1240. def test_get_invitation_token_key(self):
  1241. """Test the _get_invitation_token_key helper method."""
  1242. # Execute test
  1243. result = RegisterService._get_invitation_token_key("test-token")
  1244. # Verify results
  1245. assert result == "member_invite:token:test-token"
  1246. def test_get_invitation_by_token_with_workspace_and_email(self, mock_redis_dependencies):
  1247. """Test _get_invitation_by_token with workspace ID and email."""
  1248. # Setup mock
  1249. mock_redis_dependencies.get.return_value = b"user-123"
  1250. # Execute test
  1251. result = RegisterService._get_invitation_by_token("token-123", "workspace-456", "test@example.com")
  1252. # Verify results
  1253. assert result is not None
  1254. assert result["account_id"] == "user-123"
  1255. assert result["email"] == "test@example.com"
  1256. assert result["workspace_id"] == "workspace-456"
  1257. def test_get_invitation_by_token_without_workspace_and_email(self, mock_redis_dependencies):
  1258. """Test _get_invitation_by_token without workspace ID and email."""
  1259. # Setup mock
  1260. invitation_data = {
  1261. "account_id": "user-123",
  1262. "email": "test@example.com",
  1263. "workspace_id": "tenant-456",
  1264. }
  1265. mock_redis_dependencies.get.return_value = json.dumps(invitation_data).encode()
  1266. # Execute test
  1267. result = RegisterService._get_invitation_by_token("token-123")
  1268. # Verify results
  1269. assert result is not None
  1270. assert result == invitation_data
  1271. def test_get_invitation_by_token_no_data(self, mock_redis_dependencies):
  1272. """Test _get_invitation_by_token with no data."""
  1273. # Setup mock
  1274. mock_redis_dependencies.get.return_value = None
  1275. # Execute test
  1276. result = RegisterService._get_invitation_by_token("token-123")
  1277. # Verify results
  1278. assert result is None