test_account_service.py 126 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034303530363037303830393040304130423043304430453046304730483049305030513052305330543055305630573058305930603061306230633064306530663067306830693070307130723073307430753076307730783079308030813082308330843085308630873088308930903091309230933094309530963097309830993100310131023103310431053106310731083109311031113112311331143115311631173118311931203121312231233124312531263127312831293130313131323133313431353136313731383139314031413142314331443145314631473148314931503151315231533154315531563157315831593160316131623163316431653166316731683169317031713172317331743175317631773178317931803181318231833184318531863187318831893190319131923193319431953196319731983199320032013202320332043205320632073208320932103211321232133214321532163217321832193220322132223223322432253226322732283229323032313232323332343235323632373238323932403241324232433244324532463247324832493250325132523253325432553256325732583259326032613262326332643265326632673268326932703271327232733274327532763277327832793280328132823283328432853286328732883289329032913292329332943295329632973298329933003301330233033304330533063307330833093310331133123313331433153316331733183319332033213322332333243325332633273328332933303331333233333334333533363337333833393340
  1. import json
  2. from hashlib import sha256
  3. from unittest.mock import patch
  4. import pytest
  5. from faker import Faker
  6. from werkzeug.exceptions import Unauthorized
  7. from configs import dify_config
  8. from controllers.console.error import AccountNotFound, NotAllowedCreateWorkspace
  9. from models.account import AccountStatus, TenantAccountJoin
  10. from services.account_service import AccountService, RegisterService, TenantService, TokenPair
  11. from services.errors.account import (
  12. AccountAlreadyInTenantError,
  13. AccountLoginError,
  14. AccountNotFoundError,
  15. AccountPasswordError,
  16. AccountRegisterError,
  17. CurrentPasswordIncorrectError,
  18. )
  19. from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError
  20. class TestAccountService:
  21. """Integration tests for AccountService using testcontainers."""
  22. @pytest.fixture
  23. def mock_external_service_dependencies(self):
  24. """Mock setup for external service dependencies."""
  25. with (
  26. patch("services.account_service.FeatureService") as mock_feature_service,
  27. patch("services.account_service.BillingService") as mock_billing_service,
  28. patch("services.account_service.PassportService") as mock_passport_service,
  29. ):
  30. # Setup default mock returns
  31. mock_feature_service.get_system_features.return_value.is_allow_register = True
  32. mock_feature_service.get_system_features.return_value.is_allow_create_workspace = True
  33. mock_feature_service.get_system_features.return_value.license.workspaces.is_available.return_value = True
  34. mock_billing_service.is_email_in_freeze.return_value = False
  35. mock_passport_service.return_value.issue.return_value = "mock_jwt_token"
  36. yield {
  37. "feature_service": mock_feature_service,
  38. "billing_service": mock_billing_service,
  39. "passport_service": mock_passport_service,
  40. }
  41. def test_create_account_and_login(self, db_session_with_containers, mock_external_service_dependencies):
  42. """
  43. Test account creation and login with correct password.
  44. """
  45. fake = Faker()
  46. email = fake.email()
  47. name = fake.name()
  48. password = fake.password(length=12)
  49. # Setup mocks
  50. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  51. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  52. account = AccountService.create_account(
  53. email=email,
  54. name=name,
  55. interface_language="en-US",
  56. password=password,
  57. )
  58. assert account.email == email
  59. assert account.status == AccountStatus.ACTIVE.value
  60. # Login with correct password
  61. logged_in = AccountService.authenticate(email, password)
  62. assert logged_in.id == account.id
  63. def test_create_account_without_password(self, db_session_with_containers, mock_external_service_dependencies):
  64. """
  65. Test account creation without password (for OAuth users).
  66. """
  67. fake = Faker()
  68. email = fake.email()
  69. name = fake.name()
  70. # Setup mocks
  71. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  72. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  73. account = AccountService.create_account(
  74. email=email,
  75. name=name,
  76. interface_language="en-US",
  77. password=None,
  78. )
  79. assert account.email == email
  80. assert account.password is None
  81. assert account.password_salt is None
  82. def test_create_account_registration_disabled(self, db_session_with_containers, mock_external_service_dependencies):
  83. """
  84. Test account creation when registration is disabled.
  85. """
  86. fake = Faker()
  87. email = fake.email()
  88. name = fake.name()
  89. # Setup mocks to disable registration
  90. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = False
  91. with pytest.raises(AccountNotFound): # AccountNotFound exception
  92. AccountService.create_account(
  93. email=email,
  94. name=name,
  95. interface_language="en-US",
  96. password=fake.password(length=12),
  97. )
  98. def test_create_account_email_in_freeze(self, db_session_with_containers, mock_external_service_dependencies):
  99. """
  100. Test account creation when email is in freeze period.
  101. """
  102. fake = Faker()
  103. email = fake.email()
  104. name = fake.name()
  105. password = fake.password(length=12)
  106. # Setup mocks
  107. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  108. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = True
  109. dify_config.BILLING_ENABLED = True
  110. with pytest.raises(AccountRegisterError):
  111. AccountService.create_account(
  112. email=email,
  113. name=name,
  114. interface_language="en-US",
  115. password=password,
  116. )
  117. dify_config.BILLING_ENABLED = False # Reset config for other tests
  118. def test_authenticate_account_not_found(self, db_session_with_containers, mock_external_service_dependencies):
  119. """
  120. Test authentication with non-existent account.
  121. """
  122. fake = Faker()
  123. email = fake.email()
  124. password = fake.password(length=12)
  125. with pytest.raises(AccountNotFoundError):
  126. AccountService.authenticate(email, password)
  127. def test_authenticate_banned_account(self, db_session_with_containers, mock_external_service_dependencies):
  128. """
  129. Test authentication with banned account.
  130. """
  131. fake = Faker()
  132. email = fake.email()
  133. name = fake.name()
  134. password = fake.password(length=12)
  135. # Setup mocks
  136. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  137. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  138. # Create account first
  139. account = AccountService.create_account(
  140. email=email,
  141. name=name,
  142. interface_language="en-US",
  143. password=password,
  144. )
  145. # Ban the account
  146. account.status = AccountStatus.BANNED.value
  147. from extensions.ext_database import db
  148. db.session.commit()
  149. with pytest.raises(AccountLoginError):
  150. AccountService.authenticate(email, password)
  151. def test_authenticate_wrong_password(self, db_session_with_containers, mock_external_service_dependencies):
  152. """
  153. Test authentication with wrong password.
  154. """
  155. fake = Faker()
  156. email = fake.email()
  157. name = fake.name()
  158. correct_password = fake.password(length=12)
  159. wrong_password = fake.password(length=12)
  160. # Setup mocks
  161. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  162. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  163. # Create account first
  164. AccountService.create_account(
  165. email=email,
  166. name=name,
  167. interface_language="en-US",
  168. password=correct_password,
  169. )
  170. with pytest.raises(AccountPasswordError):
  171. AccountService.authenticate(email, wrong_password)
  172. def test_authenticate_with_invite_token(self, db_session_with_containers, mock_external_service_dependencies):
  173. """
  174. Test authentication with invite token to set password for account without password.
  175. """
  176. fake = Faker()
  177. email = fake.email()
  178. name = fake.name()
  179. new_password = fake.password(length=12)
  180. # Setup mocks
  181. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  182. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  183. # Create account without password
  184. account = AccountService.create_account(
  185. email=email,
  186. name=name,
  187. interface_language="en-US",
  188. password=None,
  189. )
  190. # Authenticate with invite token to set password
  191. authenticated_account = AccountService.authenticate(
  192. email,
  193. new_password,
  194. invite_token="valid_invite_token",
  195. )
  196. assert authenticated_account.id == account.id
  197. assert authenticated_account.password is not None
  198. assert authenticated_account.password_salt is not None
  199. def test_authenticate_pending_account_activation(
  200. self, db_session_with_containers, mock_external_service_dependencies
  201. ):
  202. """
  203. Test authentication activates pending account.
  204. """
  205. fake = Faker()
  206. email = fake.email()
  207. name = fake.name()
  208. password = fake.password(length=12)
  209. # Setup mocks
  210. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  211. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  212. # Create account with pending status
  213. account = AccountService.create_account(
  214. email=email,
  215. name=name,
  216. interface_language="en-US",
  217. password=password,
  218. )
  219. account.status = AccountStatus.PENDING.value
  220. from extensions.ext_database import db
  221. db.session.commit()
  222. # Authenticate should activate the account
  223. authenticated_account = AccountService.authenticate(email, password)
  224. assert authenticated_account.status == AccountStatus.ACTIVE.value
  225. assert authenticated_account.initialized_at is not None
  226. def test_update_account_password_success(self, db_session_with_containers, mock_external_service_dependencies):
  227. """
  228. Test successful password update.
  229. """
  230. fake = Faker()
  231. email = fake.email()
  232. name = fake.name()
  233. old_password = fake.password(length=12)
  234. new_password = fake.password(length=12)
  235. # Setup mocks
  236. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  237. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  238. # Create account
  239. account = AccountService.create_account(
  240. email=email,
  241. name=name,
  242. interface_language="en-US",
  243. password=old_password,
  244. )
  245. # Update password
  246. updated_account = AccountService.update_account_password(account, old_password, new_password)
  247. # Verify new password works
  248. authenticated_account = AccountService.authenticate(email, new_password)
  249. assert authenticated_account.id == account.id
  250. def test_update_account_password_wrong_current_password(
  251. self, db_session_with_containers, mock_external_service_dependencies
  252. ):
  253. """
  254. Test password update with wrong current password.
  255. """
  256. fake = Faker()
  257. email = fake.email()
  258. name = fake.name()
  259. old_password = fake.password(length=12)
  260. wrong_password = fake.password(length=12)
  261. new_password = fake.password(length=12)
  262. # Setup mocks
  263. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  264. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  265. # Create account
  266. account = AccountService.create_account(
  267. email=email,
  268. name=name,
  269. interface_language="en-US",
  270. password=old_password,
  271. )
  272. with pytest.raises(CurrentPasswordIncorrectError):
  273. AccountService.update_account_password(account, wrong_password, new_password)
  274. def test_update_account_password_invalid_new_password(
  275. self, db_session_with_containers, mock_external_service_dependencies
  276. ):
  277. """
  278. Test password update with invalid new password format.
  279. """
  280. fake = Faker()
  281. email = fake.email()
  282. name = fake.name()
  283. old_password = fake.password(length=12)
  284. # Setup mocks
  285. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  286. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  287. # Create account
  288. account = AccountService.create_account(
  289. email=email,
  290. name=name,
  291. interface_language="en-US",
  292. password=old_password,
  293. )
  294. # Test with too short password (assuming minimum length validation)
  295. with pytest.raises(ValueError): # Password validation error
  296. AccountService.update_account_password(account, old_password, "123")
  297. def test_create_account_and_tenant(self, db_session_with_containers, mock_external_service_dependencies):
  298. """
  299. Test account creation with automatic tenant creation.
  300. """
  301. fake = Faker()
  302. email = fake.email()
  303. name = fake.name()
  304. password = fake.password(length=12)
  305. # Setup mocks
  306. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  307. mock_external_service_dependencies[
  308. "feature_service"
  309. ].get_system_features.return_value.is_allow_create_workspace = True
  310. mock_external_service_dependencies[
  311. "feature_service"
  312. ].get_system_features.return_value.license.workspaces.is_available.return_value = True
  313. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  314. account = AccountService.create_account_and_tenant(
  315. email=email,
  316. name=name,
  317. interface_language="en-US",
  318. password=password,
  319. )
  320. assert account.email == email
  321. # Verify tenant was created and linked
  322. from extensions.ext_database import db
  323. tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first()
  324. assert tenant_join is not None
  325. assert tenant_join.role == "owner"
  326. def test_create_account_and_tenant_workspace_creation_disabled(
  327. self, db_session_with_containers, mock_external_service_dependencies
  328. ):
  329. """
  330. Test account creation when workspace creation is disabled.
  331. """
  332. fake = Faker()
  333. email = fake.email()
  334. name = fake.name()
  335. password = fake.password(length=12)
  336. # Setup mocks
  337. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  338. mock_external_service_dependencies[
  339. "feature_service"
  340. ].get_system_features.return_value.is_allow_create_workspace = False
  341. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  342. with pytest.raises(WorkSpaceNotAllowedCreateError):
  343. AccountService.create_account_and_tenant(
  344. email=email,
  345. name=name,
  346. interface_language="en-US",
  347. password=password,
  348. )
  349. def test_create_account_and_tenant_workspace_limit_exceeded(
  350. self, db_session_with_containers, mock_external_service_dependencies
  351. ):
  352. """
  353. Test account creation when workspace limit is exceeded.
  354. """
  355. fake = Faker()
  356. email = fake.email()
  357. name = fake.name()
  358. password = fake.password(length=12)
  359. # Setup mocks
  360. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  361. mock_external_service_dependencies[
  362. "feature_service"
  363. ].get_system_features.return_value.is_allow_create_workspace = True
  364. mock_external_service_dependencies[
  365. "feature_service"
  366. ].get_system_features.return_value.license.workspaces.is_available.return_value = False
  367. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  368. with pytest.raises(WorkspacesLimitExceededError):
  369. AccountService.create_account_and_tenant(
  370. email=email,
  371. name=name,
  372. interface_language="en-US",
  373. password=password,
  374. )
  375. def test_link_account_integrate_new_provider(self, db_session_with_containers, mock_external_service_dependencies):
  376. """
  377. Test linking account with new OAuth provider.
  378. """
  379. fake = Faker()
  380. email = fake.email()
  381. name = fake.name()
  382. # Setup mocks
  383. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  384. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  385. # Create account
  386. account = AccountService.create_account(
  387. email=email,
  388. name=name,
  389. interface_language="en-US",
  390. password=None,
  391. )
  392. # Link with new provider
  393. AccountService.link_account_integrate("new-google", "google_open_id_123", account)
  394. # Verify integration was created
  395. from extensions.ext_database import db
  396. from models.account import AccountIntegrate
  397. integration = db.session.query(AccountIntegrate).filter_by(account_id=account.id, provider="new-google").first()
  398. assert integration is not None
  399. assert integration.open_id == "google_open_id_123"
  400. def test_link_account_integrate_existing_provider(
  401. self, db_session_with_containers, mock_external_service_dependencies
  402. ):
  403. """
  404. Test linking account with existing provider (should update).
  405. """
  406. fake = Faker()
  407. email = fake.email()
  408. name = fake.name()
  409. # Setup mocks
  410. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  411. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  412. # Create account
  413. account = AccountService.create_account(
  414. email=email,
  415. name=name,
  416. interface_language="en-US",
  417. password=None,
  418. )
  419. # Link with provider first time
  420. AccountService.link_account_integrate("exists-google", "google_open_id_123", account)
  421. # Link with same provider but different open_id (should update)
  422. AccountService.link_account_integrate("exists-google", "google_open_id_456", account)
  423. # Verify integration was updated
  424. from extensions.ext_database import db
  425. from models.account import AccountIntegrate
  426. integration = (
  427. db.session.query(AccountIntegrate).filter_by(account_id=account.id, provider="exists-google").first()
  428. )
  429. assert integration.open_id == "google_open_id_456"
  430. def test_close_account(self, db_session_with_containers, mock_external_service_dependencies):
  431. """
  432. Test closing an account.
  433. """
  434. fake = Faker()
  435. email = fake.email()
  436. name = fake.name()
  437. password = fake.password(length=12)
  438. # Setup mocks
  439. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  440. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  441. # Create account
  442. account = AccountService.create_account(
  443. email=email,
  444. name=name,
  445. interface_language="en-US",
  446. password=password,
  447. )
  448. # Close account
  449. AccountService.close_account(account)
  450. # Verify account status changed
  451. from extensions.ext_database import db
  452. db.session.refresh(account)
  453. assert account.status == AccountStatus.CLOSED.value
  454. def test_update_account_fields(self, db_session_with_containers, mock_external_service_dependencies):
  455. """
  456. Test updating account fields.
  457. """
  458. fake = Faker()
  459. email = fake.email()
  460. name = fake.name()
  461. password = fake.password(length=12)
  462. updated_name = fake.name()
  463. # Setup mocks
  464. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  465. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  466. # Create account
  467. account = AccountService.create_account(
  468. email=email,
  469. name=name,
  470. interface_language="en-US",
  471. password=password,
  472. )
  473. # Update account fields
  474. updated_account = AccountService.update_account(account, name=updated_name, interface_theme="dark")
  475. assert updated_account.name == updated_name
  476. assert updated_account.interface_theme == "dark"
  477. def test_update_account_invalid_field(self, db_session_with_containers, mock_external_service_dependencies):
  478. """
  479. Test updating account with invalid field.
  480. """
  481. fake = Faker()
  482. email = fake.email()
  483. name = fake.name()
  484. password = fake.password(length=12)
  485. # Setup mocks
  486. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  487. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  488. # Create account
  489. account = AccountService.create_account(
  490. email=email,
  491. name=name,
  492. interface_language="en-US",
  493. password=password,
  494. )
  495. with pytest.raises(AttributeError):
  496. AccountService.update_account(account, invalid_field="value")
  497. def test_update_login_info(self, db_session_with_containers, mock_external_service_dependencies):
  498. """
  499. Test updating login information.
  500. """
  501. fake = Faker()
  502. email = fake.email()
  503. name = fake.name()
  504. password = fake.password(length=12)
  505. ip_address = fake.ipv4()
  506. # Setup mocks
  507. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  508. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  509. # Create account
  510. account = AccountService.create_account(
  511. email=email,
  512. name=name,
  513. interface_language="en-US",
  514. password=password,
  515. )
  516. # Update login info
  517. AccountService.update_login_info(account, ip_address=ip_address)
  518. # Verify login info was updated
  519. from extensions.ext_database import db
  520. db.session.refresh(account)
  521. assert account.last_login_ip == ip_address
  522. assert account.last_login_at is not None
  523. def test_login_success(self, db_session_with_containers, mock_external_service_dependencies):
  524. """
  525. Test successful login with token generation.
  526. """
  527. fake = Faker()
  528. email = fake.email()
  529. name = fake.name()
  530. password = fake.password(length=12)
  531. ip_address = fake.ipv4()
  532. # Setup mocks
  533. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  534. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  535. mock_external_service_dependencies["passport_service"].return_value.issue.return_value = "mock_access_token"
  536. # Create account
  537. account = AccountService.create_account(
  538. email=email,
  539. name=name,
  540. interface_language="en-US",
  541. password=password,
  542. )
  543. # Login
  544. token_pair = AccountService.login(account, ip_address=ip_address)
  545. assert isinstance(token_pair, TokenPair)
  546. assert token_pair.access_token == "mock_access_token"
  547. assert token_pair.refresh_token is not None
  548. # Verify passport service was called with correct parameters
  549. mock_passport = mock_external_service_dependencies["passport_service"].return_value
  550. mock_passport.issue.assert_called_once()
  551. call_args = mock_passport.issue.call_args[0][0]
  552. assert call_args["user_id"] == account.id
  553. assert call_args["iss"] is not None
  554. assert call_args["sub"] == "Console API Passport"
  555. def test_login_pending_account_activation(self, db_session_with_containers, mock_external_service_dependencies):
  556. """
  557. Test login activates pending account.
  558. """
  559. fake = Faker()
  560. email = fake.email()
  561. name = fake.name()
  562. password = fake.password(length=12)
  563. # Setup mocks
  564. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  565. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  566. mock_external_service_dependencies["passport_service"].return_value.issue.return_value = "mock_access_token"
  567. # Create account with pending status
  568. account = AccountService.create_account(
  569. email=email,
  570. name=name,
  571. interface_language="en-US",
  572. password=password,
  573. )
  574. account.status = AccountStatus.PENDING.value
  575. from extensions.ext_database import db
  576. db.session.commit()
  577. # Login should activate the account
  578. token_pair = AccountService.login(account)
  579. db.session.refresh(account)
  580. assert account.status == AccountStatus.ACTIVE.value
  581. def test_logout(self, db_session_with_containers, mock_external_service_dependencies):
  582. """
  583. Test logout functionality.
  584. """
  585. fake = Faker()
  586. email = fake.email()
  587. name = fake.name()
  588. password = fake.password(length=12)
  589. # Setup mocks
  590. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  591. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  592. mock_external_service_dependencies["passport_service"].return_value.issue.return_value = "mock_access_token"
  593. # Create account
  594. account = AccountService.create_account(
  595. email=email,
  596. name=name,
  597. interface_language="en-US",
  598. password=password,
  599. )
  600. # Login first to get refresh token
  601. token_pair = AccountService.login(account)
  602. # Logout
  603. AccountService.logout(account=account)
  604. # Verify refresh token was deleted from Redis
  605. from extensions.ext_redis import redis_client
  606. refresh_token_key = f"account_refresh_token:{account.id}"
  607. assert redis_client.get(refresh_token_key) is None
  608. def test_refresh_token_success(self, db_session_with_containers, mock_external_service_dependencies):
  609. """
  610. Test successful token refresh.
  611. """
  612. fake = Faker()
  613. email = fake.email()
  614. name = fake.name()
  615. password = fake.password(length=12)
  616. tenant_name = fake.company()
  617. # Setup mocks
  618. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  619. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  620. mock_external_service_dependencies["passport_service"].return_value.issue.return_value = "new_mock_access_token"
  621. # Create account
  622. account = AccountService.create_account(
  623. email=email,
  624. name=name,
  625. interface_language="en-US",
  626. password=password,
  627. )
  628. # Create associated Tenant
  629. TenantService.create_owner_tenant_if_not_exist(account=account, name=tenant_name, is_setup=True)
  630. # Login to get initial tokens
  631. initial_token_pair = AccountService.login(account)
  632. # Refresh token
  633. new_token_pair = AccountService.refresh_token(initial_token_pair.refresh_token)
  634. assert isinstance(new_token_pair, TokenPair)
  635. assert new_token_pair.access_token == "new_mock_access_token"
  636. assert new_token_pair.refresh_token != initial_token_pair.refresh_token
  637. def test_refresh_token_invalid_token(self, db_session_with_containers, mock_external_service_dependencies):
  638. """
  639. Test refresh token with invalid token.
  640. """
  641. fake = Faker()
  642. invalid_token = fake.uuid4()
  643. with pytest.raises(ValueError, match="Invalid refresh token"):
  644. AccountService.refresh_token(invalid_token)
  645. def test_refresh_token_invalid_account(self, db_session_with_containers, mock_external_service_dependencies):
  646. """
  647. Test refresh token with valid token but invalid account.
  648. """
  649. fake = Faker()
  650. email = fake.email()
  651. name = fake.name()
  652. password = fake.password(length=12)
  653. # Setup mocks
  654. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  655. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  656. mock_external_service_dependencies["passport_service"].return_value.issue.return_value = "mock_access_token"
  657. # Create account
  658. account = AccountService.create_account(
  659. email=email,
  660. name=name,
  661. interface_language="en-US",
  662. password=password,
  663. )
  664. # Login to get tokens
  665. token_pair = AccountService.login(account)
  666. # Delete account
  667. from extensions.ext_database import db
  668. db.session.delete(account)
  669. db.session.commit()
  670. # Try to refresh token with deleted account
  671. with pytest.raises(ValueError, match="Invalid account"):
  672. AccountService.refresh_token(token_pair.refresh_token)
  673. def test_load_user_success(self, db_session_with_containers, mock_external_service_dependencies):
  674. """
  675. Test loading user by ID successfully.
  676. """
  677. fake = Faker()
  678. email = fake.email()
  679. name = fake.name()
  680. password = fake.password(length=12)
  681. tenant_name = fake.company()
  682. # Setup mocks
  683. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  684. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  685. # Create account
  686. account = AccountService.create_account(
  687. email=email,
  688. name=name,
  689. interface_language="en-US",
  690. password=password,
  691. )
  692. # Create associated Tenant
  693. TenantService.create_owner_tenant_if_not_exist(account=account, name=tenant_name, is_setup=True)
  694. # Load user
  695. loaded_user = AccountService.load_user(account.id)
  696. assert loaded_user is not None
  697. assert loaded_user.id == account.id
  698. assert loaded_user.email == account.email
  699. def test_load_user_not_found(self, db_session_with_containers, mock_external_service_dependencies):
  700. """
  701. Test loading non-existent user.
  702. """
  703. fake = Faker()
  704. non_existent_user_id = fake.uuid4()
  705. loaded_user = AccountService.load_user(non_existent_user_id)
  706. assert loaded_user is None
  707. def test_load_user_banned_account(self, db_session_with_containers, mock_external_service_dependencies):
  708. """
  709. Test loading banned user raises Unauthorized.
  710. """
  711. fake = Faker()
  712. email = fake.email()
  713. name = fake.name()
  714. password = fake.password(length=12)
  715. # Setup mocks
  716. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  717. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  718. # Create account
  719. account = AccountService.create_account(
  720. email=email,
  721. name=name,
  722. interface_language="en-US",
  723. password=password,
  724. )
  725. # Ban the account
  726. account.status = AccountStatus.BANNED.value
  727. from extensions.ext_database import db
  728. db.session.commit()
  729. with pytest.raises(Unauthorized): # Unauthorized exception
  730. AccountService.load_user(account.id)
  731. def test_get_account_jwt_token(self, db_session_with_containers, mock_external_service_dependencies):
  732. """
  733. Test JWT token generation for account.
  734. """
  735. fake = Faker()
  736. email = fake.email()
  737. name = fake.name()
  738. password = fake.password(length=12)
  739. # Setup mocks
  740. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  741. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  742. mock_external_service_dependencies["passport_service"].return_value.issue.return_value = "mock_jwt_token"
  743. # Create account
  744. account = AccountService.create_account(
  745. email=email,
  746. name=name,
  747. interface_language="en-US",
  748. password=password,
  749. )
  750. # Generate JWT token
  751. token = AccountService.get_account_jwt_token(account)
  752. assert token == "mock_jwt_token"
  753. # Verify passport service was called with correct parameters
  754. mock_passport = mock_external_service_dependencies["passport_service"].return_value
  755. mock_passport.issue.assert_called_once()
  756. call_args = mock_passport.issue.call_args[0][0]
  757. assert call_args["user_id"] == account.id
  758. assert call_args["iss"] is not None
  759. assert call_args["sub"] == "Console API Passport"
  760. def test_load_logged_in_account(self, db_session_with_containers, mock_external_service_dependencies):
  761. """
  762. Test loading logged in account by ID.
  763. """
  764. fake = Faker()
  765. email = fake.email()
  766. name = fake.name()
  767. password = fake.password(length=12)
  768. tenant_name = fake.company()
  769. # Setup mocks
  770. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  771. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  772. # Create account
  773. account = AccountService.create_account(
  774. email=email,
  775. name=name,
  776. interface_language="en-US",
  777. password=password,
  778. )
  779. # Create associated Tenant
  780. TenantService.create_owner_tenant_if_not_exist(account=account, name=tenant_name, is_setup=True)
  781. # Load logged in account
  782. loaded_account = AccountService.load_logged_in_account(account_id=account.id)
  783. assert loaded_account is not None
  784. assert loaded_account.id == account.id
  785. def test_get_user_through_email_success(self, db_session_with_containers, mock_external_service_dependencies):
  786. """
  787. Test getting user through email successfully.
  788. """
  789. fake = Faker()
  790. email = fake.email()
  791. name = fake.name()
  792. password = fake.password(length=12)
  793. # Setup mocks
  794. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  795. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  796. # Create account
  797. account = AccountService.create_account(
  798. email=email,
  799. name=name,
  800. interface_language="en-US",
  801. password=password,
  802. )
  803. # Get user through email
  804. found_user = AccountService.get_user_through_email(email)
  805. assert found_user is not None
  806. assert found_user.id == account.id
  807. def test_get_user_through_email_not_found(self, db_session_with_containers, mock_external_service_dependencies):
  808. """
  809. Test getting user through non-existent email.
  810. """
  811. fake = Faker()
  812. non_existent_email = fake.email()
  813. found_user = AccountService.get_user_through_email(non_existent_email)
  814. assert found_user is None
  815. def test_get_user_through_email_banned_account(
  816. self, db_session_with_containers, mock_external_service_dependencies
  817. ):
  818. """
  819. Test getting banned user through email raises Unauthorized.
  820. """
  821. fake = Faker()
  822. email = fake.email()
  823. name = fake.name()
  824. password = fake.password(length=12)
  825. # Setup mocks
  826. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  827. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  828. # Create account
  829. account = AccountService.create_account(
  830. email=email,
  831. name=name,
  832. interface_language="en-US",
  833. password=password,
  834. )
  835. # Ban the account
  836. account.status = AccountStatus.BANNED.value
  837. from extensions.ext_database import db
  838. db.session.commit()
  839. with pytest.raises(Unauthorized): # Unauthorized exception
  840. AccountService.get_user_through_email(email)
  841. def test_get_user_through_email_in_freeze(self, db_session_with_containers, mock_external_service_dependencies):
  842. """
  843. Test getting user through email that is in freeze period.
  844. """
  845. fake = Faker()
  846. email_in_freeze = fake.email()
  847. # Setup mocks
  848. dify_config.BILLING_ENABLED = True
  849. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = True
  850. with pytest.raises(AccountRegisterError):
  851. AccountService.get_user_through_email(email_in_freeze)
  852. # Reset config
  853. dify_config.BILLING_ENABLED = False
  854. def test_delete_account(self, db_session_with_containers, mock_external_service_dependencies):
  855. """
  856. Test account deletion (should add task to queue).
  857. """
  858. fake = Faker()
  859. email = fake.email()
  860. name = fake.name()
  861. password = fake.password(length=12)
  862. # Setup mocks
  863. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  864. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  865. # Create account
  866. account = AccountService.create_account(
  867. email=email,
  868. name=name,
  869. interface_language="en-US",
  870. password=password,
  871. )
  872. with patch("services.account_service.delete_account_task") as mock_delete_task:
  873. # Delete account
  874. AccountService.delete_account(account)
  875. # Verify task was added to queue
  876. mock_delete_task.delay.assert_called_once_with(account.id)
  877. def test_generate_account_deletion_verification_code(
  878. self, db_session_with_containers, mock_external_service_dependencies
  879. ):
  880. """
  881. Test generating account deletion verification code.
  882. """
  883. fake = Faker()
  884. email = fake.email()
  885. name = fake.name()
  886. password = fake.password(length=12)
  887. # Setup mocks
  888. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  889. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  890. # Create account
  891. account = AccountService.create_account(
  892. email=email,
  893. name=name,
  894. interface_language="en-US",
  895. password=password,
  896. )
  897. # Generate verification code
  898. token, code = AccountService.generate_account_deletion_verification_code(account)
  899. assert token is not None
  900. assert code is not None
  901. assert len(code) == 6
  902. assert code.isdigit()
  903. def test_verify_account_deletion_code_valid(self, db_session_with_containers, mock_external_service_dependencies):
  904. """
  905. Test verifying valid account deletion code.
  906. """
  907. fake = Faker()
  908. email = fake.email()
  909. name = fake.name()
  910. password = fake.password(length=12)
  911. # Setup mocks
  912. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  913. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  914. # Create account
  915. account = AccountService.create_account(
  916. email=email,
  917. name=name,
  918. interface_language="en-US",
  919. password=password,
  920. )
  921. # Generate verification code
  922. token, code = AccountService.generate_account_deletion_verification_code(account)
  923. # Verify code
  924. is_valid = AccountService.verify_account_deletion_code(token, code)
  925. assert is_valid is True
  926. def test_verify_account_deletion_code_invalid(self, db_session_with_containers, mock_external_service_dependencies):
  927. """
  928. Test verifying invalid account deletion code.
  929. """
  930. fake = Faker()
  931. email = fake.email()
  932. name = fake.name()
  933. password = fake.password(length=12)
  934. wrong_code = fake.numerify(text="######")
  935. # Setup mocks
  936. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  937. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  938. # Create account
  939. account = AccountService.create_account(
  940. email=email,
  941. name=name,
  942. interface_language="en-US",
  943. password=password,
  944. )
  945. # Generate verification code
  946. token, code = AccountService.generate_account_deletion_verification_code(account)
  947. # Verify with wrong code
  948. is_valid = AccountService.verify_account_deletion_code(token, wrong_code)
  949. assert is_valid is False
  950. def test_verify_account_deletion_code_invalid_token(
  951. self, db_session_with_containers, mock_external_service_dependencies
  952. ):
  953. """
  954. Test verifying account deletion code with invalid token.
  955. """
  956. fake = Faker()
  957. invalid_token = fake.uuid4()
  958. invalid_code = fake.numerify(text="######")
  959. is_valid = AccountService.verify_account_deletion_code(invalid_token, invalid_code)
  960. assert is_valid is False
  961. class TestTenantService:
  962. """Integration tests for TenantService using testcontainers."""
  963. @pytest.fixture
  964. def mock_external_service_dependencies(self):
  965. """Mock setup for external service dependencies."""
  966. with (
  967. patch("services.account_service.FeatureService") as mock_feature_service,
  968. patch("services.account_service.BillingService") as mock_billing_service,
  969. ):
  970. # Setup default mock returns
  971. mock_feature_service.get_system_features.return_value.is_allow_create_workspace = True
  972. mock_feature_service.get_system_features.return_value.license.workspaces.is_available.return_value = True
  973. mock_billing_service.is_email_in_freeze.return_value = False
  974. yield {
  975. "feature_service": mock_feature_service,
  976. "billing_service": mock_billing_service,
  977. }
  978. def test_create_tenant_success(self, db_session_with_containers, mock_external_service_dependencies):
  979. """
  980. Test successful tenant creation with default settings.
  981. """
  982. fake = Faker()
  983. tenant_name = fake.company()
  984. # Setup mocks
  985. mock_external_service_dependencies[
  986. "feature_service"
  987. ].get_system_features.return_value.is_allow_create_workspace = True
  988. # Create tenant
  989. tenant = TenantService.create_tenant(name=tenant_name)
  990. assert tenant.name == tenant_name
  991. assert tenant.plan == "basic"
  992. assert tenant.status == "normal"
  993. assert tenant.encrypt_public_key is not None
  994. def test_create_tenant_workspace_creation_disabled(
  995. self, db_session_with_containers, mock_external_service_dependencies
  996. ):
  997. """
  998. Test tenant creation when workspace creation is disabled.
  999. """
  1000. fake = Faker()
  1001. tenant_name = fake.company()
  1002. # Setup mocks to disable workspace creation
  1003. mock_external_service_dependencies[
  1004. "feature_service"
  1005. ].get_system_features.return_value.is_allow_create_workspace = False
  1006. with pytest.raises(NotAllowedCreateWorkspace): # NotAllowedCreateWorkspace exception
  1007. TenantService.create_tenant(name=tenant_name)
  1008. def test_create_tenant_with_custom_name(self, db_session_with_containers, mock_external_service_dependencies):
  1009. """
  1010. Test tenant creation with custom name and setup flag.
  1011. """
  1012. fake = Faker()
  1013. custom_tenant_name = fake.company()
  1014. # Setup mocks
  1015. mock_external_service_dependencies[
  1016. "feature_service"
  1017. ].get_system_features.return_value.is_allow_create_workspace = False
  1018. # Create tenant with setup flag (should bypass workspace creation restriction)
  1019. tenant = TenantService.create_tenant(name=custom_tenant_name, is_setup=True, is_from_dashboard=True)
  1020. assert tenant.name == custom_tenant_name
  1021. assert tenant.plan == "basic"
  1022. assert tenant.status == "normal"
  1023. assert tenant.encrypt_public_key is not None
  1024. def test_create_tenant_member_success(self, db_session_with_containers, mock_external_service_dependencies):
  1025. """
  1026. Test successful tenant member creation.
  1027. """
  1028. fake = Faker()
  1029. tenant_name = fake.company()
  1030. email = fake.email()
  1031. name = fake.name()
  1032. password = fake.password(length=12)
  1033. # Setup mocks
  1034. mock_external_service_dependencies[
  1035. "feature_service"
  1036. ].get_system_features.return_value.is_allow_create_workspace = True
  1037. # Create tenant and account
  1038. tenant = TenantService.create_tenant(name=tenant_name)
  1039. account = AccountService.create_account(
  1040. email=email,
  1041. name=name,
  1042. interface_language="en-US",
  1043. password=password,
  1044. )
  1045. # Create tenant member
  1046. tenant_member = TenantService.create_tenant_member(tenant, account, role="admin")
  1047. assert tenant_member.tenant_id == tenant.id
  1048. assert tenant_member.account_id == account.id
  1049. assert tenant_member.role == "admin"
  1050. def test_create_tenant_member_duplicate_owner(self, db_session_with_containers, mock_external_service_dependencies):
  1051. """
  1052. Test creating duplicate owner for a tenant (should fail).
  1053. """
  1054. fake = Faker()
  1055. tenant_name = fake.company()
  1056. email1 = fake.email()
  1057. name1 = fake.name()
  1058. password1 = fake.password(length=12)
  1059. email2 = fake.email()
  1060. name2 = fake.name()
  1061. password2 = fake.password(length=12)
  1062. # Setup mocks
  1063. mock_external_service_dependencies[
  1064. "feature_service"
  1065. ].get_system_features.return_value.is_allow_create_workspace = True
  1066. # Create tenant and accounts
  1067. tenant = TenantService.create_tenant(name=tenant_name)
  1068. account1 = AccountService.create_account(
  1069. email=email1,
  1070. name=name1,
  1071. interface_language="en-US",
  1072. password=password1,
  1073. )
  1074. account2 = AccountService.create_account(
  1075. email=email2,
  1076. name=name2,
  1077. interface_language="en-US",
  1078. password=password2,
  1079. )
  1080. # Create first owner
  1081. TenantService.create_tenant_member(tenant, account1, role="owner")
  1082. # Try to create second owner (should fail)
  1083. with pytest.raises(Exception, match="Tenant already has an owner"):
  1084. TenantService.create_tenant_member(tenant, account2, role="owner")
  1085. def test_create_tenant_member_existing_member(self, db_session_with_containers, mock_external_service_dependencies):
  1086. """
  1087. Test updating role for existing tenant member.
  1088. """
  1089. fake = Faker()
  1090. tenant_name = fake.company()
  1091. email = fake.email()
  1092. name = fake.name()
  1093. password = fake.password(length=12)
  1094. # Setup mocks
  1095. mock_external_service_dependencies[
  1096. "feature_service"
  1097. ].get_system_features.return_value.is_allow_create_workspace = True
  1098. # Create tenant and account
  1099. tenant = TenantService.create_tenant(name=tenant_name)
  1100. account = AccountService.create_account(
  1101. email=email,
  1102. name=name,
  1103. interface_language="en-US",
  1104. password=password,
  1105. )
  1106. # Create member with initial role
  1107. tenant_member1 = TenantService.create_tenant_member(tenant, account, role="normal")
  1108. assert tenant_member1.role == "normal"
  1109. # Update member role
  1110. tenant_member2 = TenantService.create_tenant_member(tenant, account, role="editor")
  1111. assert tenant_member2.tenant_id == tenant_member1.tenant_id
  1112. assert tenant_member2.account_id == tenant_member1.account_id
  1113. assert tenant_member2.role == "editor"
  1114. def test_get_join_tenants_success(self, db_session_with_containers, mock_external_service_dependencies):
  1115. """
  1116. Test getting join tenants for an account.
  1117. """
  1118. fake = Faker()
  1119. email = fake.email()
  1120. name = fake.name()
  1121. password = fake.password(length=12)
  1122. tenant1_name = fake.company()
  1123. tenant2_name = fake.company()
  1124. # Setup mocks
  1125. mock_external_service_dependencies[
  1126. "feature_service"
  1127. ].get_system_features.return_value.is_allow_create_workspace = True
  1128. # Create account and tenants
  1129. account = AccountService.create_account(
  1130. email=email,
  1131. name=name,
  1132. interface_language="en-US",
  1133. password=password,
  1134. )
  1135. tenant1 = TenantService.create_tenant(name=tenant1_name)
  1136. tenant2 = TenantService.create_tenant(name=tenant2_name)
  1137. # Add account to both tenants
  1138. TenantService.create_tenant_member(tenant1, account, role="normal")
  1139. TenantService.create_tenant_member(tenant2, account, role="admin")
  1140. # Get join tenants
  1141. join_tenants = TenantService.get_join_tenants(account)
  1142. assert len(join_tenants) == 2
  1143. tenant_names = [tenant.name for tenant in join_tenants]
  1144. assert tenant1_name in tenant_names
  1145. assert tenant2_name in tenant_names
  1146. def test_get_current_tenant_by_account_success(
  1147. self, db_session_with_containers, mock_external_service_dependencies
  1148. ):
  1149. """
  1150. Test getting current tenant by account successfully.
  1151. """
  1152. fake = Faker()
  1153. email = fake.email()
  1154. name = fake.name()
  1155. password = fake.password(length=12)
  1156. tenant_name = fake.company()
  1157. # Setup mocks
  1158. mock_external_service_dependencies[
  1159. "feature_service"
  1160. ].get_system_features.return_value.is_allow_create_workspace = True
  1161. # Create account and tenant
  1162. account = AccountService.create_account(
  1163. email=email,
  1164. name=name,
  1165. interface_language="en-US",
  1166. password=password,
  1167. )
  1168. tenant = TenantService.create_tenant(name=tenant_name)
  1169. # Add account to tenant and set as current
  1170. TenantService.create_tenant_member(tenant, account, role="owner")
  1171. account.current_tenant = tenant
  1172. from extensions.ext_database import db
  1173. db.session.commit()
  1174. # Get current tenant
  1175. current_tenant = TenantService.get_current_tenant_by_account(account)
  1176. assert current_tenant.id == tenant.id
  1177. assert current_tenant.name == tenant.name
  1178. assert current_tenant.role == "owner"
  1179. def test_get_current_tenant_by_account_not_found(
  1180. self, db_session_with_containers, mock_external_service_dependencies
  1181. ):
  1182. """
  1183. Test getting current tenant when account has no current tenant.
  1184. """
  1185. fake = Faker()
  1186. email = fake.email()
  1187. name = fake.name()
  1188. password = fake.password(length=12)
  1189. # Setup mocks
  1190. mock_external_service_dependencies[
  1191. "feature_service"
  1192. ].get_system_features.return_value.is_allow_create_workspace = True
  1193. # Create account without setting current tenant
  1194. account = AccountService.create_account(
  1195. email=email,
  1196. name=name,
  1197. interface_language="en-US",
  1198. password=password,
  1199. )
  1200. # Try to get current tenant (should fail)
  1201. with pytest.raises(AttributeError):
  1202. TenantService.get_current_tenant_by_account(account)
  1203. def test_switch_tenant_success(self, db_session_with_containers, mock_external_service_dependencies):
  1204. """
  1205. Test successful tenant switching.
  1206. """
  1207. fake = Faker()
  1208. email = fake.email()
  1209. name = fake.name()
  1210. password = fake.password(length=12)
  1211. tenant1_name = fake.company()
  1212. tenant2_name = fake.company()
  1213. # Setup mocks
  1214. mock_external_service_dependencies[
  1215. "feature_service"
  1216. ].get_system_features.return_value.is_allow_create_workspace = True
  1217. # Create account and tenants
  1218. account = AccountService.create_account(
  1219. email=email,
  1220. name=name,
  1221. interface_language="en-US",
  1222. password=password,
  1223. )
  1224. tenant1 = TenantService.create_tenant(name=tenant1_name)
  1225. tenant2 = TenantService.create_tenant(name=tenant2_name)
  1226. # Add account to both tenants
  1227. TenantService.create_tenant_member(tenant1, account, role="owner")
  1228. TenantService.create_tenant_member(tenant2, account, role="admin")
  1229. # Set initial current tenant
  1230. account.current_tenant = tenant1
  1231. from extensions.ext_database import db
  1232. db.session.commit()
  1233. # Switch to second tenant
  1234. TenantService.switch_tenant(account, tenant2.id)
  1235. # Verify tenant was switched
  1236. db.session.refresh(account)
  1237. assert account.current_tenant_id == tenant2.id
  1238. def test_switch_tenant_no_tenant_id(self, db_session_with_containers, mock_external_service_dependencies):
  1239. """
  1240. Test tenant switching without providing tenant ID.
  1241. """
  1242. fake = Faker()
  1243. email = fake.email()
  1244. name = fake.name()
  1245. password = fake.password(length=12)
  1246. # Setup mocks
  1247. mock_external_service_dependencies[
  1248. "feature_service"
  1249. ].get_system_features.return_value.is_allow_create_workspace = True
  1250. # Create account
  1251. account = AccountService.create_account(
  1252. email=email,
  1253. name=name,
  1254. interface_language="en-US",
  1255. password=password,
  1256. )
  1257. # Try to switch tenant without providing tenant ID
  1258. with pytest.raises(ValueError, match="Tenant ID must be provided"):
  1259. TenantService.switch_tenant(account, None)
  1260. def test_switch_tenant_account_not_member(self, db_session_with_containers, mock_external_service_dependencies):
  1261. """
  1262. Test switching to a tenant where account is not a member.
  1263. """
  1264. fake = Faker()
  1265. email = fake.email()
  1266. name = fake.name()
  1267. password = fake.password(length=12)
  1268. tenant_name = fake.company()
  1269. # Setup mocks
  1270. mock_external_service_dependencies[
  1271. "feature_service"
  1272. ].get_system_features.return_value.is_allow_create_workspace = True
  1273. # Create account and tenant
  1274. account = AccountService.create_account(
  1275. email=email,
  1276. name=name,
  1277. interface_language="en-US",
  1278. password=password,
  1279. )
  1280. tenant = TenantService.create_tenant(name=tenant_name)
  1281. # Try to switch to tenant where account is not a member
  1282. with pytest.raises(Exception, match="Tenant not found or account is not a member of the tenant"):
  1283. TenantService.switch_tenant(account, tenant.id)
  1284. def test_has_roles_success(self, db_session_with_containers, mock_external_service_dependencies):
  1285. """
  1286. Test checking if tenant has specific roles.
  1287. """
  1288. fake = Faker()
  1289. tenant_name = fake.company()
  1290. owner_email = fake.email()
  1291. owner_name = fake.name()
  1292. owner_password = fake.password(length=12)
  1293. admin_email = fake.email()
  1294. admin_name = fake.name()
  1295. admin_password = fake.password(length=12)
  1296. # Setup mocks
  1297. mock_external_service_dependencies[
  1298. "feature_service"
  1299. ].get_system_features.return_value.is_allow_create_workspace = True
  1300. # Create tenant and accounts
  1301. tenant = TenantService.create_tenant(name=tenant_name)
  1302. owner_account = AccountService.create_account(
  1303. email=owner_email,
  1304. name=owner_name,
  1305. interface_language="en-US",
  1306. password=owner_password,
  1307. )
  1308. admin_account = AccountService.create_account(
  1309. email=admin_email,
  1310. name=admin_name,
  1311. interface_language="en-US",
  1312. password=admin_password,
  1313. )
  1314. # Add members with different roles
  1315. TenantService.create_tenant_member(tenant, owner_account, role="owner")
  1316. TenantService.create_tenant_member(tenant, admin_account, role="admin")
  1317. # Check if tenant has owner role
  1318. from models.account import TenantAccountRole
  1319. has_owner = TenantService.has_roles(tenant, [TenantAccountRole.OWNER])
  1320. assert has_owner is True
  1321. # Check if tenant has admin role
  1322. has_admin = TenantService.has_roles(tenant, [TenantAccountRole.ADMIN])
  1323. assert has_admin is True
  1324. # Check if tenant has normal role (should be False)
  1325. has_normal = TenantService.has_roles(tenant, [TenantAccountRole.NORMAL])
  1326. assert has_normal is False
  1327. def test_has_roles_invalid_role_type(self, db_session_with_containers, mock_external_service_dependencies):
  1328. """
  1329. Test checking roles with invalid role type.
  1330. """
  1331. fake = Faker()
  1332. tenant_name = fake.company()
  1333. invalid_role = fake.word()
  1334. # Setup mocks
  1335. mock_external_service_dependencies[
  1336. "feature_service"
  1337. ].get_system_features.return_value.is_allow_create_workspace = True
  1338. # Create tenant
  1339. tenant = TenantService.create_tenant(name=tenant_name)
  1340. # Try to check roles with invalid role type
  1341. with pytest.raises(ValueError, match="all roles must be TenantAccountRole"):
  1342. TenantService.has_roles(tenant, [invalid_role])
  1343. def test_get_user_role_success(self, db_session_with_containers, mock_external_service_dependencies):
  1344. """
  1345. Test getting user role in a tenant.
  1346. """
  1347. fake = Faker()
  1348. tenant_name = fake.company()
  1349. email = fake.email()
  1350. name = fake.name()
  1351. password = fake.password(length=12)
  1352. # Setup mocks
  1353. mock_external_service_dependencies[
  1354. "feature_service"
  1355. ].get_system_features.return_value.is_allow_create_workspace = True
  1356. # Create tenant and account
  1357. tenant = TenantService.create_tenant(name=tenant_name)
  1358. account = AccountService.create_account(
  1359. email=email,
  1360. name=name,
  1361. interface_language="en-US",
  1362. password=password,
  1363. )
  1364. # Add account to tenant with specific role
  1365. TenantService.create_tenant_member(tenant, account, role="editor")
  1366. # Get user role
  1367. user_role = TenantService.get_user_role(account, tenant)
  1368. assert user_role == "editor"
  1369. def test_check_member_permission_success(self, db_session_with_containers, mock_external_service_dependencies):
  1370. """
  1371. Test checking member permission successfully.
  1372. """
  1373. fake = Faker()
  1374. tenant_name = fake.company()
  1375. owner_email = fake.email()
  1376. owner_name = fake.name()
  1377. owner_password = fake.password(length=12)
  1378. member_email = fake.email()
  1379. member_name = fake.name()
  1380. member_password = fake.password(length=12)
  1381. # Setup mocks
  1382. mock_external_service_dependencies[
  1383. "feature_service"
  1384. ].get_system_features.return_value.is_allow_create_workspace = True
  1385. # Create tenant and accounts
  1386. tenant = TenantService.create_tenant(name=tenant_name)
  1387. owner_account = AccountService.create_account(
  1388. email=owner_email,
  1389. name=owner_name,
  1390. interface_language="en-US",
  1391. password=owner_password,
  1392. )
  1393. member_account = AccountService.create_account(
  1394. email=member_email,
  1395. name=member_name,
  1396. interface_language="en-US",
  1397. password=member_password,
  1398. )
  1399. # Add members with different roles
  1400. TenantService.create_tenant_member(tenant, owner_account, role="owner")
  1401. TenantService.create_tenant_member(tenant, member_account, role="normal")
  1402. # Check owner permission to add member (should succeed)
  1403. TenantService.check_member_permission(tenant, owner_account, member_account, "add")
  1404. def test_check_member_permission_invalid_action(
  1405. self, db_session_with_containers, mock_external_service_dependencies
  1406. ):
  1407. """
  1408. Test checking member permission with invalid action.
  1409. """
  1410. fake = Faker()
  1411. tenant_name = fake.company()
  1412. email = fake.email()
  1413. name = fake.name()
  1414. password = fake.password(length=12)
  1415. invalid_action = "invalid_action_that_doesnt_exist"
  1416. # Setup mocks
  1417. mock_external_service_dependencies[
  1418. "feature_service"
  1419. ].get_system_features.return_value.is_allow_create_workspace = True
  1420. # Create tenant and account
  1421. tenant = TenantService.create_tenant(name=tenant_name)
  1422. account = AccountService.create_account(
  1423. email=email,
  1424. name=name,
  1425. interface_language="en-US",
  1426. password=password,
  1427. )
  1428. # Add account to tenant
  1429. TenantService.create_tenant_member(tenant, account, role="owner")
  1430. # Try to check permission with invalid action
  1431. with pytest.raises(Exception, match="Invalid action"):
  1432. TenantService.check_member_permission(tenant, account, None, invalid_action)
  1433. def test_check_member_permission_operate_self(self, db_session_with_containers, mock_external_service_dependencies):
  1434. """
  1435. Test checking member permission when trying to operate self.
  1436. """
  1437. fake = Faker()
  1438. tenant_name = fake.company()
  1439. email = fake.email()
  1440. name = fake.name()
  1441. password = fake.password(length=12)
  1442. # Setup mocks
  1443. mock_external_service_dependencies[
  1444. "feature_service"
  1445. ].get_system_features.return_value.is_allow_create_workspace = True
  1446. # Create tenant and account
  1447. tenant = TenantService.create_tenant(name=tenant_name)
  1448. account = AccountService.create_account(
  1449. email=email,
  1450. name=name,
  1451. interface_language="en-US",
  1452. password=password,
  1453. )
  1454. # Add account to tenant
  1455. TenantService.create_tenant_member(tenant, account, role="owner")
  1456. # Try to check permission to operate self
  1457. with pytest.raises(Exception, match="Cannot operate self"):
  1458. TenantService.check_member_permission(tenant, account, account, "remove")
  1459. def test_remove_member_from_tenant_success(self, db_session_with_containers, mock_external_service_dependencies):
  1460. """
  1461. Test successful member removal from tenant.
  1462. """
  1463. fake = Faker()
  1464. tenant_name = fake.company()
  1465. owner_email = fake.email()
  1466. owner_name = fake.name()
  1467. owner_password = fake.password(length=12)
  1468. member_email = fake.email()
  1469. member_name = fake.name()
  1470. member_password = fake.password(length=12)
  1471. # Setup mocks
  1472. mock_external_service_dependencies[
  1473. "feature_service"
  1474. ].get_system_features.return_value.is_allow_create_workspace = True
  1475. # Create tenant and accounts
  1476. tenant = TenantService.create_tenant(name=tenant_name)
  1477. owner_account = AccountService.create_account(
  1478. email=owner_email,
  1479. name=owner_name,
  1480. interface_language="en-US",
  1481. password=owner_password,
  1482. )
  1483. member_account = AccountService.create_account(
  1484. email=member_email,
  1485. name=member_name,
  1486. interface_language="en-US",
  1487. password=member_password,
  1488. )
  1489. # Add members with different roles
  1490. TenantService.create_tenant_member(tenant, owner_account, role="owner")
  1491. TenantService.create_tenant_member(tenant, member_account, role="normal")
  1492. # Remove member
  1493. TenantService.remove_member_from_tenant(tenant, member_account, owner_account)
  1494. # Verify member was removed
  1495. from extensions.ext_database import db
  1496. from models.account import TenantAccountJoin
  1497. member_join = (
  1498. db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=member_account.id).first()
  1499. )
  1500. assert member_join is None
  1501. def test_remove_member_from_tenant_operate_self(
  1502. self, db_session_with_containers, mock_external_service_dependencies
  1503. ):
  1504. """
  1505. Test removing member when trying to operate self.
  1506. """
  1507. fake = Faker()
  1508. tenant_name = fake.company()
  1509. email = fake.email()
  1510. name = fake.name()
  1511. password = fake.password(length=12)
  1512. # Setup mocks
  1513. mock_external_service_dependencies[
  1514. "feature_service"
  1515. ].get_system_features.return_value.is_allow_create_workspace = True
  1516. # Create tenant and account
  1517. tenant = TenantService.create_tenant(name=tenant_name)
  1518. account = AccountService.create_account(
  1519. email=email,
  1520. name=name,
  1521. interface_language="en-US",
  1522. password=password,
  1523. )
  1524. # Add account to tenant
  1525. TenantService.create_tenant_member(tenant, account, role="owner")
  1526. # Try to remove self
  1527. with pytest.raises(Exception, match="Cannot operate self"):
  1528. TenantService.remove_member_from_tenant(tenant, account, account)
  1529. def test_remove_member_from_tenant_not_member(self, db_session_with_containers, mock_external_service_dependencies):
  1530. """
  1531. Test removing member who is not in the tenant.
  1532. """
  1533. fake = Faker()
  1534. tenant_name = fake.company()
  1535. owner_email = fake.email()
  1536. owner_name = fake.name()
  1537. owner_password = fake.password(length=12)
  1538. non_member_email = fake.email()
  1539. non_member_name = fake.name()
  1540. non_member_password = fake.password(length=12)
  1541. # Setup mocks
  1542. mock_external_service_dependencies[
  1543. "feature_service"
  1544. ].get_system_features.return_value.is_allow_create_workspace = True
  1545. # Create tenant and accounts
  1546. tenant = TenantService.create_tenant(name=tenant_name)
  1547. owner_account = AccountService.create_account(
  1548. email=owner_email,
  1549. name=owner_name,
  1550. interface_language="en-US",
  1551. password=owner_password,
  1552. )
  1553. non_member_account = AccountService.create_account(
  1554. email=non_member_email,
  1555. name=non_member_name,
  1556. interface_language="en-US",
  1557. password=non_member_password,
  1558. )
  1559. # Add only owner to tenant
  1560. TenantService.create_tenant_member(tenant, owner_account, role="owner")
  1561. # Try to remove non-member
  1562. with pytest.raises(Exception, match="Member not in tenant"):
  1563. TenantService.remove_member_from_tenant(tenant, non_member_account, owner_account)
  1564. def test_update_member_role_success(self, db_session_with_containers, mock_external_service_dependencies):
  1565. """
  1566. Test successful member role update.
  1567. """
  1568. fake = Faker()
  1569. tenant_name = fake.company()
  1570. owner_email = fake.email()
  1571. owner_name = fake.name()
  1572. owner_password = fake.password(length=12)
  1573. member_email = fake.email()
  1574. member_name = fake.name()
  1575. member_password = fake.password(length=12)
  1576. # Setup mocks
  1577. mock_external_service_dependencies[
  1578. "feature_service"
  1579. ].get_system_features.return_value.is_allow_create_workspace = True
  1580. # Create tenant and accounts
  1581. tenant = TenantService.create_tenant(name=tenant_name)
  1582. owner_account = AccountService.create_account(
  1583. email=owner_email,
  1584. name=owner_name,
  1585. interface_language="en-US",
  1586. password=owner_password,
  1587. )
  1588. member_account = AccountService.create_account(
  1589. email=member_email,
  1590. name=member_name,
  1591. interface_language="en-US",
  1592. password=member_password,
  1593. )
  1594. # Add members with different roles
  1595. TenantService.create_tenant_member(tenant, owner_account, role="owner")
  1596. TenantService.create_tenant_member(tenant, member_account, role="normal")
  1597. # Update member role
  1598. TenantService.update_member_role(tenant, member_account, "admin", owner_account)
  1599. # Verify role was updated
  1600. from extensions.ext_database import db
  1601. from models.account import TenantAccountJoin
  1602. member_join = (
  1603. db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=member_account.id).first()
  1604. )
  1605. assert member_join.role == "admin"
  1606. def test_update_member_role_to_owner(self, db_session_with_containers, mock_external_service_dependencies):
  1607. """
  1608. Test updating member role to owner (should change current owner to admin).
  1609. """
  1610. fake = Faker()
  1611. tenant_name = fake.company()
  1612. owner_email = fake.email()
  1613. owner_name = fake.name()
  1614. owner_password = fake.password(length=12)
  1615. member_email = fake.email()
  1616. member_name = fake.name()
  1617. member_password = fake.password(length=12)
  1618. # Setup mocks
  1619. mock_external_service_dependencies[
  1620. "feature_service"
  1621. ].get_system_features.return_value.is_allow_create_workspace = True
  1622. # Create tenant and accounts
  1623. tenant = TenantService.create_tenant(name=tenant_name)
  1624. owner_account = AccountService.create_account(
  1625. email=owner_email,
  1626. name=owner_name,
  1627. interface_language="en-US",
  1628. password=owner_password,
  1629. )
  1630. member_account = AccountService.create_account(
  1631. email=member_email,
  1632. name=member_name,
  1633. interface_language="en-US",
  1634. password=member_password,
  1635. )
  1636. # Add members with different roles
  1637. TenantService.create_tenant_member(tenant, owner_account, role="owner")
  1638. TenantService.create_tenant_member(tenant, member_account, role="admin")
  1639. # Update member role to owner
  1640. TenantService.update_member_role(tenant, member_account, "owner", owner_account)
  1641. # Verify roles were updated correctly
  1642. from extensions.ext_database import db
  1643. from models.account import TenantAccountJoin
  1644. owner_join = (
  1645. db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=owner_account.id).first()
  1646. )
  1647. member_join = (
  1648. db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=member_account.id).first()
  1649. )
  1650. assert owner_join.role == "admin"
  1651. assert member_join.role == "owner"
  1652. def test_update_member_role_already_assigned(self, db_session_with_containers, mock_external_service_dependencies):
  1653. """
  1654. Test updating member role to already assigned role.
  1655. """
  1656. fake = Faker()
  1657. tenant_name = fake.company()
  1658. owner_email = fake.email()
  1659. owner_name = fake.name()
  1660. owner_password = fake.password(length=12)
  1661. member_email = fake.email()
  1662. member_name = fake.name()
  1663. member_password = fake.password(length=12)
  1664. # Setup mocks
  1665. mock_external_service_dependencies[
  1666. "feature_service"
  1667. ].get_system_features.return_value.is_allow_create_workspace = True
  1668. # Create tenant and accounts
  1669. tenant = TenantService.create_tenant(name=tenant_name)
  1670. owner_account = AccountService.create_account(
  1671. email=owner_email,
  1672. name=owner_name,
  1673. interface_language="en-US",
  1674. password=owner_password,
  1675. )
  1676. member_account = AccountService.create_account(
  1677. email=member_email,
  1678. name=member_name,
  1679. interface_language="en-US",
  1680. password=member_password,
  1681. )
  1682. # Add members with different roles
  1683. TenantService.create_tenant_member(tenant, owner_account, role="owner")
  1684. TenantService.create_tenant_member(tenant, member_account, role="admin")
  1685. # Try to update member role to already assigned role
  1686. with pytest.raises(Exception, match="The provided role is already assigned to the member"):
  1687. TenantService.update_member_role(tenant, member_account, "admin", owner_account)
  1688. def test_get_tenant_count_success(self, db_session_with_containers, mock_external_service_dependencies):
  1689. """
  1690. Test getting tenant count successfully.
  1691. """
  1692. fake = Faker()
  1693. tenant1_name = fake.company()
  1694. tenant2_name = fake.company()
  1695. tenant3_name = fake.company()
  1696. # Setup mocks
  1697. mock_external_service_dependencies[
  1698. "feature_service"
  1699. ].get_system_features.return_value.is_allow_create_workspace = True
  1700. # Create multiple tenants
  1701. tenant1 = TenantService.create_tenant(name=tenant1_name)
  1702. tenant2 = TenantService.create_tenant(name=tenant2_name)
  1703. tenant3 = TenantService.create_tenant(name=tenant3_name)
  1704. # Get tenant count
  1705. tenant_count = TenantService.get_tenant_count()
  1706. # Should have at least 3 tenants (may be more from other tests)
  1707. assert tenant_count >= 3
  1708. def test_create_owner_tenant_if_not_exist_new_user(
  1709. self, db_session_with_containers, mock_external_service_dependencies
  1710. ):
  1711. """
  1712. Test creating owner tenant for new user without existing tenants.
  1713. """
  1714. fake = Faker()
  1715. email = fake.email()
  1716. name = fake.name()
  1717. password = fake.password(length=12)
  1718. workspace_name = fake.company()
  1719. # Setup mocks
  1720. mock_external_service_dependencies[
  1721. "feature_service"
  1722. ].get_system_features.return_value.is_allow_create_workspace = True
  1723. mock_external_service_dependencies[
  1724. "feature_service"
  1725. ].get_system_features.return_value.license.workspaces.is_available.return_value = True
  1726. # Create account
  1727. account = AccountService.create_account(
  1728. email=email,
  1729. name=name,
  1730. interface_language="en-US",
  1731. password=password,
  1732. )
  1733. # Create owner tenant
  1734. TenantService.create_owner_tenant_if_not_exist(account, name=workspace_name)
  1735. # Verify tenant was created and linked
  1736. from extensions.ext_database import db
  1737. from models.account import TenantAccountJoin
  1738. tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first()
  1739. assert tenant_join is not None
  1740. assert tenant_join.role == "owner"
  1741. assert account.current_tenant is not None
  1742. assert account.current_tenant.name == workspace_name
  1743. def test_create_owner_tenant_if_not_exist_existing_tenant(
  1744. self, db_session_with_containers, mock_external_service_dependencies
  1745. ):
  1746. """
  1747. Test creating owner tenant when user already has a tenant.
  1748. """
  1749. fake = Faker()
  1750. email = fake.email()
  1751. name = fake.name()
  1752. password = fake.password(length=12)
  1753. existing_tenant_name = fake.company()
  1754. new_workspace_name = fake.company()
  1755. # Setup mocks
  1756. mock_external_service_dependencies[
  1757. "feature_service"
  1758. ].get_system_features.return_value.is_allow_create_workspace = True
  1759. mock_external_service_dependencies[
  1760. "feature_service"
  1761. ].get_system_features.return_value.license.workspaces.is_available.return_value = True
  1762. # Create account and existing tenant
  1763. account = AccountService.create_account(
  1764. email=email,
  1765. name=name,
  1766. interface_language="en-US",
  1767. password=password,
  1768. )
  1769. existing_tenant = TenantService.create_tenant(name=existing_tenant_name)
  1770. TenantService.create_tenant_member(existing_tenant, account, role="owner")
  1771. account.current_tenant = existing_tenant
  1772. from extensions.ext_database import db
  1773. db.session.commit()
  1774. # Try to create owner tenant again (should not create new one)
  1775. TenantService.create_owner_tenant_if_not_exist(account, name=new_workspace_name)
  1776. # Verify no new tenant was created
  1777. tenant_joins = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).all()
  1778. assert len(tenant_joins) == 1
  1779. assert account.current_tenant.id == existing_tenant.id
  1780. def test_create_owner_tenant_if_not_exist_workspace_disabled(
  1781. self, db_session_with_containers, mock_external_service_dependencies
  1782. ):
  1783. """
  1784. Test creating owner tenant when workspace creation is disabled.
  1785. """
  1786. fake = Faker()
  1787. email = fake.email()
  1788. name = fake.name()
  1789. password = fake.password(length=12)
  1790. workspace_name = fake.company()
  1791. # Setup mocks to disable workspace creation
  1792. mock_external_service_dependencies[
  1793. "feature_service"
  1794. ].get_system_features.return_value.is_allow_create_workspace = False
  1795. # Create account
  1796. account = AccountService.create_account(
  1797. email=email,
  1798. name=name,
  1799. interface_language="en-US",
  1800. password=password,
  1801. )
  1802. # Try to create owner tenant (should fail)
  1803. with pytest.raises(WorkSpaceNotAllowedCreateError): # WorkSpaceNotAllowedCreateError exception
  1804. TenantService.create_owner_tenant_if_not_exist(account, name=workspace_name)
  1805. def test_get_tenant_members_success(self, db_session_with_containers, mock_external_service_dependencies):
  1806. """
  1807. Test getting tenant members successfully.
  1808. """
  1809. fake = Faker()
  1810. tenant_name = fake.company()
  1811. owner_email = fake.email()
  1812. owner_name = fake.name()
  1813. owner_password = fake.password(length=12)
  1814. admin_email = fake.email()
  1815. admin_name = fake.name()
  1816. admin_password = fake.password(length=12)
  1817. normal_email = fake.email()
  1818. normal_name = fake.name()
  1819. normal_password = fake.password(length=12)
  1820. # Setup mocks
  1821. mock_external_service_dependencies[
  1822. "feature_service"
  1823. ].get_system_features.return_value.is_allow_create_workspace = True
  1824. # Create tenant and accounts
  1825. tenant = TenantService.create_tenant(name=tenant_name)
  1826. owner_account = AccountService.create_account(
  1827. email=owner_email,
  1828. name=owner_name,
  1829. interface_language="en-US",
  1830. password=owner_password,
  1831. )
  1832. admin_account = AccountService.create_account(
  1833. email=admin_email,
  1834. name=admin_name,
  1835. interface_language="en-US",
  1836. password=admin_password,
  1837. )
  1838. normal_account = AccountService.create_account(
  1839. email=normal_email,
  1840. name=normal_name,
  1841. interface_language="en-US",
  1842. password=normal_password,
  1843. )
  1844. # Add members with different roles
  1845. TenantService.create_tenant_member(tenant, owner_account, role="owner")
  1846. TenantService.create_tenant_member(tenant, admin_account, role="admin")
  1847. TenantService.create_tenant_member(tenant, normal_account, role="normal")
  1848. # Get tenant members
  1849. members = TenantService.get_tenant_members(tenant)
  1850. assert len(members) == 3
  1851. member_emails = [member.email for member in members]
  1852. assert owner_email in member_emails
  1853. assert admin_email in member_emails
  1854. assert normal_email in member_emails
  1855. # Verify roles are set correctly
  1856. for member in members:
  1857. if member.email == owner_email:
  1858. assert member.role == "owner"
  1859. elif member.email == admin_email:
  1860. assert member.role == "admin"
  1861. elif member.email == normal_email:
  1862. assert member.role == "normal"
  1863. def test_get_dataset_operator_members_success(self, db_session_with_containers, mock_external_service_dependencies):
  1864. """
  1865. Test getting dataset operator members successfully.
  1866. """
  1867. fake = Faker()
  1868. tenant_name = fake.company()
  1869. owner_email = fake.email()
  1870. owner_name = fake.name()
  1871. owner_password = fake.password(length=12)
  1872. operator_email = fake.email()
  1873. operator_name = fake.name()
  1874. operator_password = fake.password(length=12)
  1875. normal_email = fake.email()
  1876. normal_name = fake.name()
  1877. normal_password = fake.password(length=12)
  1878. # Setup mocks
  1879. mock_external_service_dependencies[
  1880. "feature_service"
  1881. ].get_system_features.return_value.is_allow_create_workspace = True
  1882. # Create tenant and accounts
  1883. tenant = TenantService.create_tenant(name=tenant_name)
  1884. owner_account = AccountService.create_account(
  1885. email=owner_email,
  1886. name=owner_name,
  1887. interface_language="en-US",
  1888. password=owner_password,
  1889. )
  1890. dataset_operator_account = AccountService.create_account(
  1891. email=operator_email,
  1892. name=operator_name,
  1893. interface_language="en-US",
  1894. password=operator_password,
  1895. )
  1896. normal_account = AccountService.create_account(
  1897. email=normal_email,
  1898. name=normal_name,
  1899. interface_language="en-US",
  1900. password=normal_password,
  1901. )
  1902. # Add members with different roles
  1903. TenantService.create_tenant_member(tenant, owner_account, role="owner")
  1904. TenantService.create_tenant_member(tenant, dataset_operator_account, role="dataset_operator")
  1905. TenantService.create_tenant_member(tenant, normal_account, role="normal")
  1906. # Get dataset operator members
  1907. dataset_operators = TenantService.get_dataset_operator_members(tenant)
  1908. assert len(dataset_operators) == 1
  1909. assert dataset_operators[0].email == operator_email
  1910. assert dataset_operators[0].role == "dataset_operator"
  1911. def test_get_custom_config_success(self, db_session_with_containers, mock_external_service_dependencies):
  1912. """
  1913. Test getting custom config successfully.
  1914. """
  1915. fake = Faker()
  1916. tenant_name = fake.company()
  1917. theme = fake.random_element(elements=("dark", "light"))
  1918. language = fake.random_element(elements=("zh-CN", "en-US"))
  1919. # Setup mocks
  1920. mock_external_service_dependencies[
  1921. "feature_service"
  1922. ].get_system_features.return_value.is_allow_create_workspace = True
  1923. # Create tenant with custom config
  1924. tenant = TenantService.create_tenant(name=tenant_name)
  1925. # Set custom config
  1926. custom_config = {"theme": theme, "language": language, "feature_flags": {"beta": True}}
  1927. tenant.custom_config_dict = custom_config
  1928. from extensions.ext_database import db
  1929. db.session.commit()
  1930. # Get custom config
  1931. retrieved_config = TenantService.get_custom_config(tenant.id)
  1932. assert retrieved_config == custom_config
  1933. assert retrieved_config["theme"] == theme
  1934. assert retrieved_config["language"] == language
  1935. assert retrieved_config["feature_flags"]["beta"] is True
  1936. class TestRegisterService:
  1937. """Integration tests for RegisterService using testcontainers."""
  1938. @pytest.fixture
  1939. def mock_external_service_dependencies(self):
  1940. """Mock setup for external service dependencies."""
  1941. with (
  1942. patch("services.account_service.FeatureService") as mock_feature_service,
  1943. patch("services.account_service.BillingService") as mock_billing_service,
  1944. patch("services.account_service.PassportService") as mock_passport_service,
  1945. ):
  1946. # Setup default mock returns
  1947. mock_feature_service.get_system_features.return_value.is_allow_register = True
  1948. mock_feature_service.get_system_features.return_value.is_allow_create_workspace = True
  1949. mock_feature_service.get_system_features.return_value.license.workspaces.is_available.return_value = True
  1950. mock_billing_service.is_email_in_freeze.return_value = False
  1951. mock_passport_service.return_value.issue.return_value = "mock_jwt_token"
  1952. yield {
  1953. "feature_service": mock_feature_service,
  1954. "billing_service": mock_billing_service,
  1955. "passport_service": mock_passport_service,
  1956. }
  1957. def test_setup_success(self, db_session_with_containers, mock_external_service_dependencies):
  1958. """
  1959. Test successful system setup with account creation and tenant setup.
  1960. """
  1961. fake = Faker()
  1962. admin_email = fake.email()
  1963. admin_name = fake.name()
  1964. admin_password = fake.password(length=12)
  1965. ip_address = fake.ipv4()
  1966. # Setup mocks
  1967. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  1968. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  1969. # Execute setup
  1970. RegisterService.setup(
  1971. email=admin_email,
  1972. name=admin_name,
  1973. password=admin_password,
  1974. ip_address=ip_address,
  1975. )
  1976. # Verify account was created
  1977. from extensions.ext_database import db
  1978. from models.account import Account
  1979. from models.model import DifySetup
  1980. account = db.session.query(Account).filter_by(email=admin_email).first()
  1981. assert account is not None
  1982. assert account.name == admin_name
  1983. assert account.last_login_ip == ip_address
  1984. assert account.initialized_at is not None
  1985. assert account.status == "active"
  1986. # Verify DifySetup was created
  1987. dify_setup = db.session.query(DifySetup).first()
  1988. assert dify_setup is not None
  1989. # Verify tenant was created and linked
  1990. from models.account import TenantAccountJoin
  1991. tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first()
  1992. assert tenant_join is not None
  1993. assert tenant_join.role == "owner"
  1994. def test_setup_failure_rollback(self, db_session_with_containers, mock_external_service_dependencies):
  1995. """
  1996. Test setup failure with proper rollback of all created entities.
  1997. """
  1998. fake = Faker()
  1999. admin_email = fake.email()
  2000. admin_name = fake.name()
  2001. admin_password = fake.password(length=12)
  2002. ip_address = fake.ipv4()
  2003. # Setup mocks
  2004. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2005. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2006. # Mock AccountService.create_account to raise exception
  2007. with patch("services.account_service.AccountService.create_account") as mock_create_account:
  2008. mock_create_account.side_effect = Exception("Database error")
  2009. # Execute setup and verify exception
  2010. with pytest.raises(ValueError, match="Setup failed: Database error"):
  2011. RegisterService.setup(
  2012. email=admin_email,
  2013. name=admin_name,
  2014. password=admin_password,
  2015. ip_address=ip_address,
  2016. )
  2017. # Verify no entities were created (rollback worked)
  2018. from extensions.ext_database import db
  2019. from models.account import Account, Tenant, TenantAccountJoin
  2020. from models.model import DifySetup
  2021. account = db.session.query(Account).filter_by(email=admin_email).first()
  2022. tenant_count = db.session.query(Tenant).count()
  2023. tenant_join_count = db.session.query(TenantAccountJoin).count()
  2024. dify_setup_count = db.session.query(DifySetup).count()
  2025. assert account is None
  2026. assert tenant_count == 0
  2027. assert tenant_join_count == 0
  2028. assert dify_setup_count == 0
  2029. def test_register_success(self, db_session_with_containers, mock_external_service_dependencies):
  2030. """
  2031. Test successful account registration with workspace creation.
  2032. """
  2033. fake = Faker()
  2034. email = fake.email()
  2035. name = fake.name()
  2036. password = fake.password(length=12)
  2037. language = fake.random_element(elements=("en-US", "zh-CN"))
  2038. # Setup mocks
  2039. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2040. mock_external_service_dependencies[
  2041. "feature_service"
  2042. ].get_system_features.return_value.is_allow_create_workspace = True
  2043. mock_external_service_dependencies[
  2044. "feature_service"
  2045. ].get_system_features.return_value.license.workspaces.is_available.return_value = True
  2046. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2047. # Execute registration
  2048. account = RegisterService.register(
  2049. email=email,
  2050. name=name,
  2051. password=password,
  2052. language=language,
  2053. )
  2054. # Verify account was created
  2055. assert account.email == email
  2056. assert account.name == name
  2057. assert account.status == "active"
  2058. assert account.initialized_at is not None
  2059. # Verify tenant was created and linked
  2060. from extensions.ext_database import db
  2061. from models.account import TenantAccountJoin
  2062. tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first()
  2063. assert tenant_join is not None
  2064. assert tenant_join.role == "owner"
  2065. assert account.current_tenant is not None
  2066. assert account.current_tenant.name == f"{name}'s Workspace"
  2067. def test_register_with_oauth(self, db_session_with_containers, mock_external_service_dependencies):
  2068. """
  2069. Test account registration with OAuth integration.
  2070. """
  2071. fake = Faker()
  2072. email = fake.email()
  2073. name = fake.name()
  2074. open_id = fake.uuid4()
  2075. provider = fake.random_element(elements=("google", "github", "microsoft"))
  2076. language = fake.random_element(elements=("en-US", "zh-CN"))
  2077. # Setup mocks
  2078. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2079. mock_external_service_dependencies[
  2080. "feature_service"
  2081. ].get_system_features.return_value.is_allow_create_workspace = True
  2082. mock_external_service_dependencies[
  2083. "feature_service"
  2084. ].get_system_features.return_value.license.workspaces.is_available.return_value = True
  2085. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2086. # Execute registration with OAuth
  2087. account = RegisterService.register(
  2088. email=email,
  2089. name=name,
  2090. password=None,
  2091. open_id=open_id,
  2092. provider=provider,
  2093. language=language,
  2094. )
  2095. # Verify account was created
  2096. assert account.email == email
  2097. assert account.name == name
  2098. assert account.status == "active"
  2099. assert account.initialized_at is not None
  2100. # Verify OAuth integration was created
  2101. from extensions.ext_database import db
  2102. from models.account import AccountIntegrate
  2103. integration = db.session.query(AccountIntegrate).filter_by(account_id=account.id, provider=provider).first()
  2104. assert integration is not None
  2105. assert integration.open_id == open_id
  2106. def test_register_with_pending_status(self, db_session_with_containers, mock_external_service_dependencies):
  2107. """
  2108. Test account registration with pending status.
  2109. """
  2110. fake = Faker()
  2111. email = fake.email()
  2112. name = fake.name()
  2113. password = fake.password(length=12)
  2114. language = fake.random_element(elements=("en-US", "zh-CN"))
  2115. # Setup mocks
  2116. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2117. mock_external_service_dependencies[
  2118. "feature_service"
  2119. ].get_system_features.return_value.is_allow_create_workspace = True
  2120. mock_external_service_dependencies[
  2121. "feature_service"
  2122. ].get_system_features.return_value.license.workspaces.is_available.return_value = True
  2123. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2124. # Execute registration with pending status
  2125. from models.account import AccountStatus
  2126. account = RegisterService.register(
  2127. email=email,
  2128. name=name,
  2129. password=password,
  2130. language=language,
  2131. status=AccountStatus.PENDING,
  2132. )
  2133. # Verify account was created with pending status
  2134. assert account.email == email
  2135. assert account.name == name
  2136. assert account.status == "pending"
  2137. assert account.initialized_at is not None
  2138. # Verify tenant was created and linked
  2139. from extensions.ext_database import db
  2140. from models.account import TenantAccountJoin
  2141. tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first()
  2142. assert tenant_join is not None
  2143. assert tenant_join.role == "owner"
  2144. def test_register_workspace_creation_disabled(self, db_session_with_containers, mock_external_service_dependencies):
  2145. """
  2146. Test account registration when workspace creation is disabled.
  2147. """
  2148. fake = Faker()
  2149. email = fake.email()
  2150. name = fake.name()
  2151. password = fake.password(length=12)
  2152. language = fake.random_element(elements=("en-US", "zh-CN"))
  2153. # Setup mocks
  2154. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2155. mock_external_service_dependencies[
  2156. "feature_service"
  2157. ].get_system_features.return_value.is_allow_create_workspace = False
  2158. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2159. # with pytest.raises(AccountRegisterError, match="Workspace is not allowed to create."):
  2160. account = RegisterService.register(
  2161. email=email,
  2162. name=name,
  2163. password=password,
  2164. language=language,
  2165. )
  2166. # Verify account was created with no tenant
  2167. assert account.email == email
  2168. assert account.name == name
  2169. assert account.status == "active"
  2170. assert account.initialized_at is not None
  2171. # Verify tenant was created and linked
  2172. from extensions.ext_database import db
  2173. from models.account import TenantAccountJoin
  2174. tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first()
  2175. assert tenant_join is None
  2176. def test_register_workspace_limit_exceeded(self, db_session_with_containers, mock_external_service_dependencies):
  2177. """
  2178. Test account registration when workspace limit is exceeded.
  2179. """
  2180. fake = Faker()
  2181. email = fake.email()
  2182. name = fake.name()
  2183. password = fake.password(length=12)
  2184. language = fake.random_element(elements=("en-US", "zh-CN"))
  2185. # Setup mocks
  2186. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2187. mock_external_service_dependencies[
  2188. "feature_service"
  2189. ].get_system_features.return_value.is_allow_create_workspace = True
  2190. mock_external_service_dependencies[
  2191. "feature_service"
  2192. ].get_system_features.return_value.license.workspaces.is_available.return_value = False
  2193. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2194. # with pytest.raises(AccountRegisterError, match="Workspace is not allowed to create."):
  2195. account = RegisterService.register(
  2196. email=email,
  2197. name=name,
  2198. password=password,
  2199. language=language,
  2200. )
  2201. # Verify account was created with no tenant
  2202. assert account.email == email
  2203. assert account.name == name
  2204. assert account.status == "active"
  2205. assert account.initialized_at is not None
  2206. # Verify tenant was created and linked
  2207. from extensions.ext_database import db
  2208. from models.account import TenantAccountJoin
  2209. tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first()
  2210. assert tenant_join is None
  2211. def test_register_without_workspace(self, db_session_with_containers, mock_external_service_dependencies):
  2212. """
  2213. Test account registration without workspace creation.
  2214. """
  2215. fake = Faker()
  2216. email = fake.email()
  2217. name = fake.name()
  2218. password = fake.password(length=12)
  2219. language = fake.random_element(elements=("en-US", "zh-CN"))
  2220. # Setup mocks
  2221. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2222. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2223. # Execute registration without workspace creation
  2224. account = RegisterService.register(
  2225. email=email,
  2226. name=name,
  2227. password=password,
  2228. language=language,
  2229. create_workspace_required=False,
  2230. )
  2231. # Verify account was created
  2232. assert account.email == email
  2233. assert account.name == name
  2234. assert account.status == "active"
  2235. assert account.initialized_at is not None
  2236. # Verify no tenant was created
  2237. from extensions.ext_database import db
  2238. from models.account import TenantAccountJoin
  2239. tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first()
  2240. assert tenant_join is None
  2241. def test_invite_new_member_new_account(self, db_session_with_containers, mock_external_service_dependencies):
  2242. """
  2243. Test inviting a new member who doesn't have an account yet.
  2244. """
  2245. fake = Faker()
  2246. tenant_name = fake.company()
  2247. inviter_email = fake.email()
  2248. inviter_name = fake.name()
  2249. inviter_password = fake.password(length=12)
  2250. new_member_email = fake.email()
  2251. language = fake.random_element(elements=("en-US", "zh-CN"))
  2252. # Setup mocks
  2253. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2254. mock_external_service_dependencies[
  2255. "feature_service"
  2256. ].get_system_features.return_value.is_allow_create_workspace = True
  2257. mock_external_service_dependencies[
  2258. "feature_service"
  2259. ].get_system_features.return_value.license.workspaces.is_available.return_value = True
  2260. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2261. # Create tenant and inviter account
  2262. tenant = TenantService.create_tenant(name=tenant_name)
  2263. inviter = AccountService.create_account(
  2264. email=inviter_email,
  2265. name=inviter_name,
  2266. interface_language="en-US",
  2267. password=inviter_password,
  2268. )
  2269. TenantService.create_tenant_member(tenant, inviter, role="owner")
  2270. # Mock the email task
  2271. with patch("services.account_service.send_invite_member_mail_task") as mock_send_mail:
  2272. mock_send_mail.delay.return_value = None
  2273. # Execute invitation
  2274. token = RegisterService.invite_new_member(
  2275. tenant=tenant,
  2276. email=new_member_email,
  2277. language=language,
  2278. role="normal",
  2279. inviter=inviter,
  2280. )
  2281. # Verify token was generated
  2282. assert token is not None
  2283. assert len(token) > 0
  2284. # Verify email task was called
  2285. mock_send_mail.delay.assert_called_once()
  2286. # Verify new account was created with pending status
  2287. from extensions.ext_database import db
  2288. from models.account import Account, TenantAccountJoin
  2289. new_account = db.session.query(Account).filter_by(email=new_member_email).first()
  2290. assert new_account is not None
  2291. assert new_account.name == new_member_email.split("@")[0] # Default name from email
  2292. assert new_account.status == "pending"
  2293. # Verify tenant member was created
  2294. tenant_join = (
  2295. db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=new_account.id).first()
  2296. )
  2297. assert tenant_join is not None
  2298. assert tenant_join.role == "normal"
  2299. def test_invite_new_member_existing_account(self, db_session_with_containers, mock_external_service_dependencies):
  2300. """
  2301. Test inviting an existing member who is not in the tenant yet.
  2302. """
  2303. fake = Faker()
  2304. tenant_name = fake.company()
  2305. inviter_email = fake.email()
  2306. inviter_name = fake.name()
  2307. inviter_password = fake.password(length=12)
  2308. existing_member_email = fake.email()
  2309. existing_member_name = fake.name()
  2310. existing_member_password = fake.password(length=12)
  2311. language = fake.random_element(elements=("en-US", "zh-CN"))
  2312. # Setup mocks
  2313. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2314. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2315. # Create tenant and inviter account
  2316. tenant = TenantService.create_tenant(name=tenant_name)
  2317. inviter = AccountService.create_account(
  2318. email=inviter_email,
  2319. name=inviter_name,
  2320. interface_language="en-US",
  2321. password=inviter_password,
  2322. )
  2323. TenantService.create_tenant_member(tenant, inviter, role="owner")
  2324. # Create existing account
  2325. existing_account = AccountService.create_account(
  2326. email=existing_member_email,
  2327. name=existing_member_name,
  2328. interface_language="en-US",
  2329. password=existing_member_password,
  2330. )
  2331. # Mock the email task
  2332. with patch("services.account_service.send_invite_member_mail_task") as mock_send_mail:
  2333. mock_send_mail.delay.return_value = None
  2334. with pytest.raises(AccountAlreadyInTenantError, match="Account already in tenant."):
  2335. # Execute invitation
  2336. token = RegisterService.invite_new_member(
  2337. tenant=tenant,
  2338. email=existing_member_email,
  2339. language=language,
  2340. role="admin",
  2341. inviter=inviter,
  2342. )
  2343. # Verify email task was not called
  2344. mock_send_mail.delay.assert_not_called()
  2345. # Verify tenant member was created for existing account
  2346. from extensions.ext_database import db
  2347. from models.account import TenantAccountJoin
  2348. tenant_join = (
  2349. db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=existing_account.id).first()
  2350. )
  2351. assert tenant_join is not None
  2352. assert tenant_join.role == "admin"
  2353. def test_invite_new_member_existing_member(self, db_session_with_containers, mock_external_service_dependencies):
  2354. """
  2355. Test inviting a member who is already in the tenant with pending status.
  2356. """
  2357. fake = Faker()
  2358. tenant_name = fake.company()
  2359. inviter_email = fake.email()
  2360. inviter_name = fake.name()
  2361. inviter_password = fake.password(length=12)
  2362. existing_pending_member_email = fake.email()
  2363. existing_pending_member_name = fake.name()
  2364. existing_pending_member_password = fake.password(length=12)
  2365. language = fake.random_element(elements=("en-US", "zh-CN"))
  2366. # Setup mocks
  2367. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2368. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2369. # Create tenant and inviter account
  2370. tenant = TenantService.create_tenant(name=tenant_name)
  2371. inviter = AccountService.create_account(
  2372. email=inviter_email,
  2373. name=inviter_name,
  2374. interface_language="en-US",
  2375. password=inviter_password,
  2376. )
  2377. TenantService.create_tenant_member(tenant, inviter, role="owner")
  2378. # Create existing account with pending status
  2379. existing_account = AccountService.create_account(
  2380. email=existing_pending_member_email,
  2381. name=existing_pending_member_name,
  2382. interface_language="en-US",
  2383. password=existing_pending_member_password,
  2384. )
  2385. existing_account.status = "pending"
  2386. from extensions.ext_database import db
  2387. db.session.commit()
  2388. # Add existing account to tenant
  2389. TenantService.create_tenant_member(tenant, existing_account, role="normal")
  2390. # Mock the email task
  2391. with patch("services.account_service.send_invite_member_mail_task") as mock_send_mail:
  2392. mock_send_mail.delay.return_value = None
  2393. # Execute invitation (should resend email for pending member)
  2394. token = RegisterService.invite_new_member(
  2395. tenant=tenant,
  2396. email=existing_pending_member_email,
  2397. language=language,
  2398. role="normal",
  2399. inviter=inviter,
  2400. )
  2401. # Verify token was generated
  2402. assert token is not None
  2403. assert len(token) > 0
  2404. # Verify email task was called
  2405. mock_send_mail.delay.assert_called_once()
  2406. def test_invite_new_member_no_inviter(self, db_session_with_containers, mock_external_service_dependencies):
  2407. """
  2408. Test inviting a member without providing an inviter.
  2409. """
  2410. fake = Faker()
  2411. tenant_name = fake.company()
  2412. new_member_email = fake.email()
  2413. language = fake.random_element(elements=("en-US", "zh-CN"))
  2414. # Setup mocks
  2415. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2416. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2417. # Create tenant
  2418. tenant = TenantService.create_tenant(name=tenant_name)
  2419. # Execute invitation without inviter (should fail)
  2420. with pytest.raises(ValueError, match="Inviter is required"):
  2421. RegisterService.invite_new_member(
  2422. tenant=tenant,
  2423. email=new_member_email,
  2424. language=language,
  2425. role="normal",
  2426. inviter=None,
  2427. )
  2428. def test_invite_new_member_account_already_in_tenant(
  2429. self, db_session_with_containers, mock_external_service_dependencies
  2430. ):
  2431. """
  2432. Test inviting a member who is already in the tenant with active status.
  2433. """
  2434. fake = Faker()
  2435. tenant_name = fake.company()
  2436. inviter_email = fake.email()
  2437. inviter_name = fake.name()
  2438. inviter_password = fake.password(length=12)
  2439. already_in_tenant_email = fake.email()
  2440. already_in_tenant_name = fake.name()
  2441. already_in_tenant_password = fake.password(length=12)
  2442. language = fake.random_element(elements=("en-US", "zh-CN"))
  2443. # Setup mocks
  2444. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2445. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2446. # Create tenant and inviter account
  2447. tenant = TenantService.create_tenant(name=tenant_name)
  2448. inviter = AccountService.create_account(
  2449. email=inviter_email,
  2450. name=inviter_name,
  2451. interface_language="en-US",
  2452. password=inviter_password,
  2453. )
  2454. TenantService.create_tenant_member(tenant, inviter, role="owner")
  2455. # Create existing account with active status
  2456. existing_account = AccountService.create_account(
  2457. email=already_in_tenant_email,
  2458. name=already_in_tenant_name,
  2459. interface_language="en-US",
  2460. password=already_in_tenant_password,
  2461. )
  2462. existing_account.status = "active"
  2463. from extensions.ext_database import db
  2464. db.session.commit()
  2465. # Add existing account to tenant
  2466. TenantService.create_tenant_member(tenant, existing_account, role="normal")
  2467. # Execute invitation (should fail for active member)
  2468. with pytest.raises(AccountAlreadyInTenantError, match="Account already in tenant."):
  2469. RegisterService.invite_new_member(
  2470. tenant=tenant,
  2471. email=already_in_tenant_email,
  2472. language=language,
  2473. role="normal",
  2474. inviter=inviter,
  2475. )
  2476. def test_generate_invite_token_success(self, db_session_with_containers, mock_external_service_dependencies):
  2477. """
  2478. Test successful generation of invite token.
  2479. """
  2480. fake = Faker()
  2481. tenant_name = fake.company()
  2482. email = fake.email()
  2483. name = fake.name()
  2484. password = fake.password(length=12)
  2485. # Setup mocks
  2486. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2487. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2488. # Create tenant and account
  2489. tenant = TenantService.create_tenant(name=tenant_name)
  2490. account = AccountService.create_account(
  2491. email=email,
  2492. name=name,
  2493. interface_language="en-US",
  2494. password=password,
  2495. )
  2496. # Execute token generation
  2497. token = RegisterService.generate_invite_token(tenant, account)
  2498. # Verify token was generated
  2499. assert token is not None
  2500. assert len(token) > 0
  2501. # Verify token was stored in Redis
  2502. from extensions.ext_redis import redis_client
  2503. token_key = RegisterService._get_invitation_token_key(token)
  2504. stored_data = redis_client.get(token_key)
  2505. assert stored_data is not None
  2506. # Verify stored data contains correct information
  2507. import json
  2508. invitation_data = json.loads(stored_data.decode("utf-8"))
  2509. assert invitation_data["account_id"] == str(account.id)
  2510. assert invitation_data["email"] == account.email
  2511. assert invitation_data["workspace_id"] == tenant.id
  2512. def test_is_valid_invite_token_valid(self, db_session_with_containers, mock_external_service_dependencies):
  2513. """
  2514. Test validation of valid invite token.
  2515. """
  2516. fake = Faker()
  2517. tenant_name = fake.company()
  2518. email = fake.email()
  2519. name = fake.name()
  2520. password = fake.password(length=12)
  2521. # Setup mocks
  2522. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2523. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2524. # Create tenant and account
  2525. tenant = TenantService.create_tenant(name=tenant_name)
  2526. account = AccountService.create_account(
  2527. email=email,
  2528. name=name,
  2529. interface_language="en-US",
  2530. password=password,
  2531. )
  2532. # Generate a real token
  2533. token = RegisterService.generate_invite_token(tenant, account)
  2534. # Execute validation
  2535. is_valid = RegisterService.is_valid_invite_token(token)
  2536. # Verify token is valid
  2537. assert is_valid is True
  2538. def test_is_valid_invite_token_invalid(self, db_session_with_containers, mock_external_service_dependencies):
  2539. """
  2540. Test validation of invalid invite token.
  2541. """
  2542. fake = Faker()
  2543. invalid_token = fake.uuid4()
  2544. # Execute validation with non-existent token
  2545. is_valid = RegisterService.is_valid_invite_token(invalid_token)
  2546. # Verify token is invalid
  2547. assert is_valid is False
  2548. def test_revoke_token_with_workspace_and_email(
  2549. self, db_session_with_containers, mock_external_service_dependencies
  2550. ):
  2551. """
  2552. Test revoking token with workspace ID and email.
  2553. """
  2554. fake = Faker()
  2555. tenant_name = fake.company()
  2556. email = fake.email()
  2557. name = fake.name()
  2558. password = fake.password(length=12)
  2559. # Setup mocks
  2560. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2561. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2562. # Create tenant and account
  2563. tenant = TenantService.create_tenant(name=tenant_name)
  2564. account = AccountService.create_account(
  2565. email=email,
  2566. name=name,
  2567. interface_language="en-US",
  2568. password=password,
  2569. )
  2570. # Generate a real token
  2571. token = RegisterService.generate_invite_token(tenant, account)
  2572. # Verify token exists in Redis before revocation
  2573. from extensions.ext_redis import redis_client
  2574. token_key = RegisterService._get_invitation_token_key(token)
  2575. assert redis_client.get(token_key) is not None
  2576. # Execute token revocation
  2577. RegisterService.revoke_token(
  2578. workspace_id=tenant.id,
  2579. email=account.email,
  2580. token=token,
  2581. )
  2582. # Verify token was not deleted from Redis
  2583. assert redis_client.get(token_key) is not None
  2584. def test_revoke_token_without_workspace_and_email(
  2585. self, db_session_with_containers, mock_external_service_dependencies
  2586. ):
  2587. """
  2588. Test revoking token without workspace ID and email.
  2589. """
  2590. fake = Faker()
  2591. tenant_name = fake.company()
  2592. email = fake.email()
  2593. name = fake.name()
  2594. password = fake.password(length=12)
  2595. # Setup mocks
  2596. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2597. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2598. # Create tenant and account
  2599. tenant = TenantService.create_tenant(name=tenant_name)
  2600. account = AccountService.create_account(
  2601. email=email,
  2602. name=name,
  2603. interface_language="en-US",
  2604. password=password,
  2605. )
  2606. # Generate a real token
  2607. token = RegisterService.generate_invite_token(tenant, account)
  2608. # Verify token exists in Redis before revocation
  2609. from extensions.ext_redis import redis_client
  2610. token_key = RegisterService._get_invitation_token_key(token)
  2611. assert redis_client.get(token_key) is not None
  2612. # Execute token revocation without workspace and email
  2613. RegisterService.revoke_token(
  2614. workspace_id="",
  2615. email="",
  2616. token=token,
  2617. )
  2618. # Verify token was deleted from Redis
  2619. assert redis_client.get(token_key) is None
  2620. def test_get_invitation_if_token_valid_success(
  2621. self, db_session_with_containers, mock_external_service_dependencies
  2622. ):
  2623. """
  2624. Test getting invitation data with valid token.
  2625. """
  2626. fake = Faker()
  2627. tenant_name = fake.company()
  2628. email = fake.email()
  2629. name = fake.name()
  2630. password = fake.password(length=12)
  2631. # Setup mocks
  2632. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2633. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2634. # Create tenant and account
  2635. tenant = TenantService.create_tenant(name=tenant_name)
  2636. account = AccountService.create_account(
  2637. email=email,
  2638. name=name,
  2639. interface_language="en-US",
  2640. password=password,
  2641. )
  2642. TenantService.create_tenant_member(tenant, account, role="normal")
  2643. # Generate a real token
  2644. token = RegisterService.generate_invite_token(tenant, account)
  2645. email_hash = sha256(account.email.encode()).hexdigest()
  2646. cache_key = f"member_invite_token:{tenant.id}, {email_hash}:{token}"
  2647. from extensions.ext_redis import redis_client
  2648. redis_client.setex(cache_key, 24 * 60 * 60, account.id)
  2649. # Execute invitation retrieval
  2650. result = RegisterService.get_invitation_if_token_valid(
  2651. workspace_id=tenant.id,
  2652. email=account.email,
  2653. token=token,
  2654. )
  2655. # Verify result contains expected data
  2656. assert result is not None
  2657. assert result["account"].id == account.id
  2658. assert result["tenant"].id == tenant.id
  2659. assert result["data"]["account_id"] == str(account.id)
  2660. assert result["data"]["email"] == account.email
  2661. assert result["data"]["workspace_id"] == tenant.id
  2662. def test_get_invitation_if_token_valid_invalid_token(
  2663. self, db_session_with_containers, mock_external_service_dependencies
  2664. ):
  2665. """
  2666. Test getting invitation data with invalid token.
  2667. """
  2668. fake = Faker()
  2669. workspace_id = fake.uuid4()
  2670. email = fake.email()
  2671. invalid_token = fake.uuid4()
  2672. # Execute invitation retrieval with invalid token
  2673. result = RegisterService.get_invitation_if_token_valid(
  2674. workspace_id=workspace_id,
  2675. email=email,
  2676. token=invalid_token,
  2677. )
  2678. # Verify result is None
  2679. assert result is None
  2680. def test_get_invitation_if_token_valid_invalid_tenant(
  2681. self, db_session_with_containers, mock_external_service_dependencies
  2682. ):
  2683. """
  2684. Test getting invitation data with invalid tenant.
  2685. """
  2686. fake = Faker()
  2687. email = fake.email()
  2688. name = fake.name()
  2689. password = fake.password(length=12)
  2690. invalid_tenant_id = fake.uuid4()
  2691. token = fake.uuid4()
  2692. # Setup mocks
  2693. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2694. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2695. # Create account
  2696. account = AccountService.create_account(
  2697. email=email,
  2698. name=name,
  2699. interface_language="en-US",
  2700. password=password,
  2701. )
  2702. # Create a real token but with non-existent tenant ID
  2703. from extensions.ext_redis import redis_client
  2704. invitation_data = {
  2705. "account_id": str(account.id),
  2706. "email": account.email,
  2707. "workspace_id": invalid_tenant_id,
  2708. }
  2709. token_key = RegisterService._get_invitation_token_key(token)
  2710. import json
  2711. redis_client.setex(token_key, 24 * 60 * 60, json.dumps(invitation_data))
  2712. # Execute invitation retrieval
  2713. result = RegisterService.get_invitation_if_token_valid(
  2714. workspace_id=invalid_tenant_id,
  2715. email=account.email,
  2716. token=token,
  2717. )
  2718. # Verify result is None (tenant not found)
  2719. assert result is None
  2720. # Clean up
  2721. redis_client.delete(token_key)
  2722. def test_get_invitation_if_token_valid_account_mismatch(
  2723. self, db_session_with_containers, mock_external_service_dependencies
  2724. ):
  2725. """
  2726. Test getting invitation data with account ID mismatch.
  2727. """
  2728. fake = Faker()
  2729. tenant_name = fake.company()
  2730. email = fake.email()
  2731. name = fake.name()
  2732. password = fake.password(length=12)
  2733. token = fake.uuid4()
  2734. # Setup mocks
  2735. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2736. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2737. # Create tenant and account
  2738. tenant = TenantService.create_tenant(name=tenant_name)
  2739. account = AccountService.create_account(
  2740. email=email,
  2741. name=name,
  2742. interface_language="en-US",
  2743. password=password,
  2744. )
  2745. TenantService.create_tenant_member(tenant, account, role="normal")
  2746. # Create a real token but with mismatched account ID
  2747. from extensions.ext_redis import redis_client
  2748. invitation_data = {
  2749. "account_id": "different-account-id", # Different from actual account ID
  2750. "email": account.email,
  2751. "workspace_id": tenant.id,
  2752. }
  2753. token_key = RegisterService._get_invitation_token_key(token)
  2754. redis_client.setex(token_key, 24 * 60 * 60, json.dumps(invitation_data))
  2755. # Execute invitation retrieval
  2756. result = RegisterService.get_invitation_if_token_valid(
  2757. workspace_id=tenant.id,
  2758. email=account.email,
  2759. token=token,
  2760. )
  2761. # Verify result is None (account ID mismatch)
  2762. assert result is None
  2763. # Clean up
  2764. redis_client.delete(token_key)
  2765. def test_get_invitation_if_token_valid_tenant_not_normal(
  2766. self, db_session_with_containers, mock_external_service_dependencies
  2767. ):
  2768. """
  2769. Test getting invitation data with tenant not in normal status.
  2770. """
  2771. fake = Faker()
  2772. tenant_name = fake.company()
  2773. email = fake.email()
  2774. name = fake.name()
  2775. password = fake.password(length=12)
  2776. token = fake.uuid4()
  2777. # Setup mocks
  2778. mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
  2779. mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
  2780. # Create tenant and account
  2781. tenant = TenantService.create_tenant(name=tenant_name)
  2782. account = AccountService.create_account(
  2783. email=email,
  2784. name=name,
  2785. interface_language="en-US",
  2786. password=password,
  2787. )
  2788. TenantService.create_tenant_member(tenant, account, role="normal")
  2789. # Change tenant status to non-normal
  2790. tenant.status = "suspended"
  2791. from extensions.ext_database import db
  2792. db.session.commit()
  2793. # Create a real token
  2794. from extensions.ext_redis import redis_client
  2795. invitation_data = {
  2796. "account_id": str(account.id),
  2797. "email": account.email,
  2798. "workspace_id": tenant.id,
  2799. }
  2800. token_key = RegisterService._get_invitation_token_key(token)
  2801. import json
  2802. redis_client.setex(token_key, 24 * 60 * 60, json.dumps(invitation_data))
  2803. # Execute invitation retrieval
  2804. result = RegisterService.get_invitation_if_token_valid(
  2805. workspace_id=tenant.id,
  2806. email=account.email,
  2807. token=token,
  2808. )
  2809. # Verify result is None (tenant not in normal status)
  2810. assert result is None
  2811. # Clean up
  2812. redis_client.delete(token_key)
  2813. def test_get_invitation_by_token_with_workspace_and_email(
  2814. self, db_session_with_containers, mock_external_service_dependencies
  2815. ):
  2816. """
  2817. Test getting invitation by token with workspace ID and email.
  2818. """
  2819. fake = Faker()
  2820. token = fake.uuid4()
  2821. workspace_id = fake.uuid4()
  2822. email = fake.email()
  2823. # Create the cache key as the service does
  2824. from hashlib import sha256
  2825. from extensions.ext_redis import redis_client
  2826. email_hash = sha256(email.encode()).hexdigest()
  2827. cache_key = f"member_invite_token:{workspace_id}, {email_hash}:{token}"
  2828. # Store account ID in Redis
  2829. account_id = fake.uuid4()
  2830. redis_client.setex(cache_key, 24 * 60 * 60, account_id)
  2831. # Execute invitation retrieval
  2832. result = RegisterService._get_invitation_by_token(
  2833. token=token,
  2834. workspace_id=workspace_id,
  2835. email=email,
  2836. )
  2837. # Verify result contains expected data
  2838. assert result is not None
  2839. assert result["account_id"] == account_id
  2840. assert result["email"] == email
  2841. assert result["workspace_id"] == workspace_id
  2842. # Clean up
  2843. redis_client.delete(cache_key)
  2844. def test_get_invitation_by_token_without_workspace_and_email(
  2845. self, db_session_with_containers, mock_external_service_dependencies
  2846. ):
  2847. """
  2848. Test getting invitation by token without workspace ID and email.
  2849. """
  2850. fake = Faker()
  2851. token = fake.uuid4()
  2852. invitation_data = {
  2853. "account_id": fake.uuid4(),
  2854. "email": fake.email(),
  2855. "workspace_id": fake.uuid4(),
  2856. }
  2857. # Store invitation data in Redis using standard token key
  2858. from extensions.ext_redis import redis_client
  2859. token_key = RegisterService._get_invitation_token_key(token)
  2860. import json
  2861. redis_client.setex(token_key, 24 * 60 * 60, json.dumps(invitation_data))
  2862. # Execute invitation retrieval
  2863. result = RegisterService._get_invitation_by_token(token=token)
  2864. # Verify result contains expected data
  2865. assert result is not None
  2866. assert result["account_id"] == invitation_data["account_id"]
  2867. assert result["email"] == invitation_data["email"]
  2868. assert result["workspace_id"] == invitation_data["workspace_id"]
  2869. # Clean up
  2870. redis_client.delete(token_key)
  2871. def test_get_invitation_token_key(self, db_session_with_containers, mock_external_service_dependencies):
  2872. """
  2873. Test getting invitation token key.
  2874. """
  2875. fake = Faker()
  2876. token = fake.uuid4()
  2877. # Execute token key generation
  2878. token_key = RegisterService._get_invitation_token_key(token)
  2879. # Verify token key format
  2880. assert token_key == f"member_invite:token:{token}"