Nevar pievienot vairāk kā 25 tēmas Tēmai ir jāsākas ar burtu vai ciparu, tā var saturēt domu zīmes ('-') un var būt līdz 35 simboliem gara.

test_account_service.py 67KB

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