| 1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034303530363037303830393040304130423043304430453046304730483049305030513052305330543055305630573058305930603061306230633064306530663067306830693070307130723073307430753076307730783079308030813082308330843085308630873088308930903091309230933094309530963097309830993100310131023103310431053106310731083109311031113112311331143115311631173118311931203121312231233124312531263127312831293130313131323133313431353136313731383139314031413142314331443145314631473148314931503151315231533154315531563157315831593160316131623163316431653166316731683169317031713172317331743175317631773178317931803181318231833184318531863187318831893190319131923193319431953196319731983199320032013202320332043205320632073208320932103211321232133214321532163217321832193220322132223223322432253226322732283229323032313232323332343235323632373238323932403241324232433244324532463247324832493250325132523253325432553256325732583259326032613262326332643265326632673268326932703271327232733274327532763277327832793280328132823283328432853286328732883289329032913292329332943295329632973298329933003301330233033304330533063307330833093310331133123313331433153316331733183319332033213322332333243325332633273328332933303331333233333334333533363337333833393340 |
- import json
- from hashlib import sha256
- from unittest.mock import patch
- import pytest
- from faker import Faker
- from werkzeug.exceptions import Unauthorized
- from configs import dify_config
- from controllers.console.error import AccountNotFound, NotAllowedCreateWorkspace
- from models.account import AccountStatus, TenantAccountJoin
- from services.account_service import AccountService, RegisterService, TenantService, TokenPair
- from services.errors.account import (
- AccountAlreadyInTenantError,
- AccountLoginError,
- AccountNotFoundError,
- AccountPasswordError,
- AccountRegisterError,
- CurrentPasswordIncorrectError,
- )
- from services.errors.workspace import WorkSpaceNotAllowedCreateError, WorkspacesLimitExceededError
- class TestAccountService:
- """Integration tests for AccountService using testcontainers."""
- @pytest.fixture
- def mock_external_service_dependencies(self):
- """Mock setup for external service dependencies."""
- with (
- patch("services.account_service.FeatureService") as mock_feature_service,
- patch("services.account_service.BillingService") as mock_billing_service,
- patch("services.account_service.PassportService") as mock_passport_service,
- ):
- # Setup default mock returns
- mock_feature_service.get_system_features.return_value.is_allow_register = True
- mock_feature_service.get_system_features.return_value.is_allow_create_workspace = True
- mock_feature_service.get_system_features.return_value.license.workspaces.is_available.return_value = True
- mock_billing_service.is_email_in_freeze.return_value = False
- mock_passport_service.return_value.issue.return_value = "mock_jwt_token"
- yield {
- "feature_service": mock_feature_service,
- "billing_service": mock_billing_service,
- "passport_service": mock_passport_service,
- }
- def test_create_account_and_login(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test account creation and login with correct password.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- assert account.email == email
- assert account.status == AccountStatus.ACTIVE.value
- # Login with correct password
- logged_in = AccountService.authenticate(email, password)
- assert logged_in.id == account.id
- def test_create_account_without_password(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test account creation without password (for OAuth users).
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=None,
- )
- assert account.email == email
- assert account.password is None
- assert account.password_salt is None
- def test_create_account_registration_disabled(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test account creation when registration is disabled.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- # Setup mocks to disable registration
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = False
- with pytest.raises(AccountNotFound): # AccountNotFound exception
- AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=fake.password(length=12),
- )
- def test_create_account_email_in_freeze(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test account creation when email is in freeze period.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = True
- dify_config.BILLING_ENABLED = True
- with pytest.raises(AccountRegisterError):
- AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- dify_config.BILLING_ENABLED = False # Reset config for other tests
- def test_authenticate_account_not_found(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test authentication with non-existent account.
- """
- fake = Faker()
- email = fake.email()
- password = fake.password(length=12)
- with pytest.raises(AccountNotFoundError):
- AccountService.authenticate(email, password)
- def test_authenticate_banned_account(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test authentication with banned account.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account first
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Ban the account
- account.status = AccountStatus.BANNED.value
- from extensions.ext_database import db
- db.session.commit()
- with pytest.raises(AccountLoginError):
- AccountService.authenticate(email, password)
- def test_authenticate_wrong_password(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test authentication with wrong password.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- correct_password = fake.password(length=12)
- wrong_password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account first
- AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=correct_password,
- )
- with pytest.raises(AccountPasswordError):
- AccountService.authenticate(email, wrong_password)
- def test_authenticate_with_invite_token(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test authentication with invite token to set password for account without password.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- new_password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account without password
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=None,
- )
- # Authenticate with invite token to set password
- authenticated_account = AccountService.authenticate(
- email,
- new_password,
- invite_token="valid_invite_token",
- )
- assert authenticated_account.id == account.id
- assert authenticated_account.password is not None
- assert authenticated_account.password_salt is not None
- def test_authenticate_pending_account_activation(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test authentication activates pending account.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account with pending status
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- account.status = AccountStatus.PENDING.value
- from extensions.ext_database import db
- db.session.commit()
- # Authenticate should activate the account
- authenticated_account = AccountService.authenticate(email, password)
- assert authenticated_account.status == AccountStatus.ACTIVE.value
- assert authenticated_account.initialized_at is not None
- def test_update_account_password_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test successful password update.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- old_password = fake.password(length=12)
- new_password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=old_password,
- )
- # Update password
- updated_account = AccountService.update_account_password(account, old_password, new_password)
- # Verify new password works
- authenticated_account = AccountService.authenticate(email, new_password)
- assert authenticated_account.id == account.id
- def test_update_account_password_wrong_current_password(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test password update with wrong current password.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- old_password = fake.password(length=12)
- wrong_password = fake.password(length=12)
- new_password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=old_password,
- )
- with pytest.raises(CurrentPasswordIncorrectError):
- AccountService.update_account_password(account, wrong_password, new_password)
- def test_update_account_password_invalid_new_password(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test password update with invalid new password format.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- old_password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=old_password,
- )
- # Test with too short password (assuming minimum length validation)
- with pytest.raises(ValueError): # Password validation error
- AccountService.update_account_password(account, old_password, "123")
- def test_create_account_and_tenant(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test account creation with automatic tenant creation.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.license.workspaces.is_available.return_value = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- account = AccountService.create_account_and_tenant(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- assert account.email == email
- # Verify tenant was created and linked
- from extensions.ext_database import db
- tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first()
- assert tenant_join is not None
- assert tenant_join.role == "owner"
- def test_create_account_and_tenant_workspace_creation_disabled(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test account creation when workspace creation is disabled.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = False
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- with pytest.raises(WorkSpaceNotAllowedCreateError):
- AccountService.create_account_and_tenant(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- def test_create_account_and_tenant_workspace_limit_exceeded(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test account creation when workspace limit is exceeded.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.license.workspaces.is_available.return_value = False
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- with pytest.raises(WorkspacesLimitExceededError):
- AccountService.create_account_and_tenant(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- def test_link_account_integrate_new_provider(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test linking account with new OAuth provider.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=None,
- )
- # Link with new provider
- AccountService.link_account_integrate("new-google", "google_open_id_123", account)
- # Verify integration was created
- from extensions.ext_database import db
- from models.account import AccountIntegrate
- integration = db.session.query(AccountIntegrate).filter_by(account_id=account.id, provider="new-google").first()
- assert integration is not None
- assert integration.open_id == "google_open_id_123"
- def test_link_account_integrate_existing_provider(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test linking account with existing provider (should update).
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=None,
- )
- # Link with provider first time
- AccountService.link_account_integrate("exists-google", "google_open_id_123", account)
- # Link with same provider but different open_id (should update)
- AccountService.link_account_integrate("exists-google", "google_open_id_456", account)
- # Verify integration was updated
- from extensions.ext_database import db
- from models.account import AccountIntegrate
- integration = (
- db.session.query(AccountIntegrate).filter_by(account_id=account.id, provider="exists-google").first()
- )
- assert integration.open_id == "google_open_id_456"
- def test_close_account(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test closing an account.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Close account
- AccountService.close_account(account)
- # Verify account status changed
- from extensions.ext_database import db
- db.session.refresh(account)
- assert account.status == AccountStatus.CLOSED.value
- def test_update_account_fields(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test updating account fields.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- updated_name = fake.name()
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Update account fields
- updated_account = AccountService.update_account(account, name=updated_name, interface_theme="dark")
- assert updated_account.name == updated_name
- assert updated_account.interface_theme == "dark"
- def test_update_account_invalid_field(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test updating account with invalid field.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- with pytest.raises(AttributeError):
- AccountService.update_account(account, invalid_field="value")
- def test_update_login_info(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test updating login information.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- ip_address = fake.ipv4()
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Update login info
- AccountService.update_login_info(account, ip_address=ip_address)
- # Verify login info was updated
- from extensions.ext_database import db
- db.session.refresh(account)
- assert account.last_login_ip == ip_address
- assert account.last_login_at is not None
- def test_login_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test successful login with token generation.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- ip_address = fake.ipv4()
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- mock_external_service_dependencies["passport_service"].return_value.issue.return_value = "mock_access_token"
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Login
- token_pair = AccountService.login(account, ip_address=ip_address)
- assert isinstance(token_pair, TokenPair)
- assert token_pair.access_token == "mock_access_token"
- assert token_pair.refresh_token is not None
- # Verify passport service was called with correct parameters
- mock_passport = mock_external_service_dependencies["passport_service"].return_value
- mock_passport.issue.assert_called_once()
- call_args = mock_passport.issue.call_args[0][0]
- assert call_args["user_id"] == account.id
- assert call_args["iss"] is not None
- assert call_args["sub"] == "Console API Passport"
- def test_login_pending_account_activation(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test login activates pending account.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- mock_external_service_dependencies["passport_service"].return_value.issue.return_value = "mock_access_token"
- # Create account with pending status
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- account.status = AccountStatus.PENDING.value
- from extensions.ext_database import db
- db.session.commit()
- # Login should activate the account
- token_pair = AccountService.login(account)
- db.session.refresh(account)
- assert account.status == AccountStatus.ACTIVE.value
- def test_logout(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test logout functionality.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- mock_external_service_dependencies["passport_service"].return_value.issue.return_value = "mock_access_token"
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Login first to get refresh token
- token_pair = AccountService.login(account)
- # Logout
- AccountService.logout(account=account)
- # Verify refresh token was deleted from Redis
- from extensions.ext_redis import redis_client
- refresh_token_key = f"account_refresh_token:{account.id}"
- assert redis_client.get(refresh_token_key) is None
- def test_refresh_token_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test successful token refresh.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- tenant_name = fake.company()
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- mock_external_service_dependencies["passport_service"].return_value.issue.return_value = "new_mock_access_token"
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Create associated Tenant
- TenantService.create_owner_tenant_if_not_exist(account=account, name=tenant_name, is_setup=True)
- # Login to get initial tokens
- initial_token_pair = AccountService.login(account)
- # Refresh token
- new_token_pair = AccountService.refresh_token(initial_token_pair.refresh_token)
- assert isinstance(new_token_pair, TokenPair)
- assert new_token_pair.access_token == "new_mock_access_token"
- assert new_token_pair.refresh_token != initial_token_pair.refresh_token
- def test_refresh_token_invalid_token(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test refresh token with invalid token.
- """
- fake = Faker()
- invalid_token = fake.uuid4()
- with pytest.raises(ValueError, match="Invalid refresh token"):
- AccountService.refresh_token(invalid_token)
- def test_refresh_token_invalid_account(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test refresh token with valid token but invalid account.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- mock_external_service_dependencies["passport_service"].return_value.issue.return_value = "mock_access_token"
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Login to get tokens
- token_pair = AccountService.login(account)
- # Delete account
- from extensions.ext_database import db
- db.session.delete(account)
- db.session.commit()
- # Try to refresh token with deleted account
- with pytest.raises(ValueError, match="Invalid account"):
- AccountService.refresh_token(token_pair.refresh_token)
- def test_load_user_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test loading user by ID successfully.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- tenant_name = fake.company()
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Create associated Tenant
- TenantService.create_owner_tenant_if_not_exist(account=account, name=tenant_name, is_setup=True)
- # Load user
- loaded_user = AccountService.load_user(account.id)
- assert loaded_user is not None
- assert loaded_user.id == account.id
- assert loaded_user.email == account.email
- def test_load_user_not_found(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test loading non-existent user.
- """
- fake = Faker()
- non_existent_user_id = fake.uuid4()
- loaded_user = AccountService.load_user(non_existent_user_id)
- assert loaded_user is None
- def test_load_user_banned_account(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test loading banned user raises Unauthorized.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Ban the account
- account.status = AccountStatus.BANNED.value
- from extensions.ext_database import db
- db.session.commit()
- with pytest.raises(Unauthorized): # Unauthorized exception
- AccountService.load_user(account.id)
- def test_get_account_jwt_token(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test JWT token generation for account.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- mock_external_service_dependencies["passport_service"].return_value.issue.return_value = "mock_jwt_token"
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Generate JWT token
- token = AccountService.get_account_jwt_token(account)
- assert token == "mock_jwt_token"
- # Verify passport service was called with correct parameters
- mock_passport = mock_external_service_dependencies["passport_service"].return_value
- mock_passport.issue.assert_called_once()
- call_args = mock_passport.issue.call_args[0][0]
- assert call_args["user_id"] == account.id
- assert call_args["iss"] is not None
- assert call_args["sub"] == "Console API Passport"
- def test_load_logged_in_account(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test loading logged in account by ID.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- tenant_name = fake.company()
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Create associated Tenant
- TenantService.create_owner_tenant_if_not_exist(account=account, name=tenant_name, is_setup=True)
- # Load logged in account
- loaded_account = AccountService.load_logged_in_account(account_id=account.id)
- assert loaded_account is not None
- assert loaded_account.id == account.id
- def test_get_user_through_email_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test getting user through email successfully.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Get user through email
- found_user = AccountService.get_user_through_email(email)
- assert found_user is not None
- assert found_user.id == account.id
- def test_get_user_through_email_not_found(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test getting user through non-existent email.
- """
- fake = Faker()
- non_existent_email = fake.email()
- found_user = AccountService.get_user_through_email(non_existent_email)
- assert found_user is None
- def test_get_user_through_email_banned_account(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test getting banned user through email raises Unauthorized.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Ban the account
- account.status = AccountStatus.BANNED.value
- from extensions.ext_database import db
- db.session.commit()
- with pytest.raises(Unauthorized): # Unauthorized exception
- AccountService.get_user_through_email(email)
- def test_get_user_through_email_in_freeze(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test getting user through email that is in freeze period.
- """
- fake = Faker()
- email_in_freeze = fake.email()
- # Setup mocks
- dify_config.BILLING_ENABLED = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = True
- with pytest.raises(AccountRegisterError):
- AccountService.get_user_through_email(email_in_freeze)
- # Reset config
- dify_config.BILLING_ENABLED = False
- def test_delete_account(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test account deletion (should add task to queue).
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- with patch("services.account_service.delete_account_task") as mock_delete_task:
- # Delete account
- AccountService.delete_account(account)
- # Verify task was added to queue
- mock_delete_task.delay.assert_called_once_with(account.id)
- def test_generate_account_deletion_verification_code(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test generating account deletion verification code.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Generate verification code
- token, code = AccountService.generate_account_deletion_verification_code(account)
- assert token is not None
- assert code is not None
- assert len(code) == 6
- assert code.isdigit()
- def test_verify_account_deletion_code_valid(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test verifying valid account deletion code.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Generate verification code
- token, code = AccountService.generate_account_deletion_verification_code(account)
- # Verify code
- is_valid = AccountService.verify_account_deletion_code(token, code)
- assert is_valid is True
- def test_verify_account_deletion_code_invalid(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test verifying invalid account deletion code.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- wrong_code = fake.numerify(text="######")
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Generate verification code
- token, code = AccountService.generate_account_deletion_verification_code(account)
- # Verify with wrong code
- is_valid = AccountService.verify_account_deletion_code(token, wrong_code)
- assert is_valid is False
- def test_verify_account_deletion_code_invalid_token(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test verifying account deletion code with invalid token.
- """
- fake = Faker()
- invalid_token = fake.uuid4()
- invalid_code = fake.numerify(text="######")
- is_valid = AccountService.verify_account_deletion_code(invalid_token, invalid_code)
- assert is_valid is False
- class TestTenantService:
- """Integration tests for TenantService using testcontainers."""
- @pytest.fixture
- def mock_external_service_dependencies(self):
- """Mock setup for external service dependencies."""
- with (
- patch("services.account_service.FeatureService") as mock_feature_service,
- patch("services.account_service.BillingService") as mock_billing_service,
- ):
- # Setup default mock returns
- mock_feature_service.get_system_features.return_value.is_allow_create_workspace = True
- mock_feature_service.get_system_features.return_value.license.workspaces.is_available.return_value = True
- mock_billing_service.is_email_in_freeze.return_value = False
- yield {
- "feature_service": mock_feature_service,
- "billing_service": mock_billing_service,
- }
- def test_create_tenant_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test successful tenant creation with default settings.
- """
- fake = Faker()
- tenant_name = fake.company()
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant
- tenant = TenantService.create_tenant(name=tenant_name)
- assert tenant.name == tenant_name
- assert tenant.plan == "basic"
- assert tenant.status == "normal"
- assert tenant.encrypt_public_key is not None
- def test_create_tenant_workspace_creation_disabled(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test tenant creation when workspace creation is disabled.
- """
- fake = Faker()
- tenant_name = fake.company()
- # Setup mocks to disable workspace creation
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = False
- with pytest.raises(NotAllowedCreateWorkspace): # NotAllowedCreateWorkspace exception
- TenantService.create_tenant(name=tenant_name)
- def test_create_tenant_with_custom_name(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test tenant creation with custom name and setup flag.
- """
- fake = Faker()
- custom_tenant_name = fake.company()
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = False
- # Create tenant with setup flag (should bypass workspace creation restriction)
- tenant = TenantService.create_tenant(name=custom_tenant_name, is_setup=True, is_from_dashboard=True)
- assert tenant.name == custom_tenant_name
- assert tenant.plan == "basic"
- assert tenant.status == "normal"
- assert tenant.encrypt_public_key is not None
- def test_create_tenant_member_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test successful tenant member creation.
- """
- fake = Faker()
- tenant_name = fake.company()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant and account
- tenant = TenantService.create_tenant(name=tenant_name)
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Create tenant member
- tenant_member = TenantService.create_tenant_member(tenant, account, role="admin")
- assert tenant_member.tenant_id == tenant.id
- assert tenant_member.account_id == account.id
- assert tenant_member.role == "admin"
- def test_create_tenant_member_duplicate_owner(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test creating duplicate owner for a tenant (should fail).
- """
- fake = Faker()
- tenant_name = fake.company()
- email1 = fake.email()
- name1 = fake.name()
- password1 = fake.password(length=12)
- email2 = fake.email()
- name2 = fake.name()
- password2 = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant and accounts
- tenant = TenantService.create_tenant(name=tenant_name)
- account1 = AccountService.create_account(
- email=email1,
- name=name1,
- interface_language="en-US",
- password=password1,
- )
- account2 = AccountService.create_account(
- email=email2,
- name=name2,
- interface_language="en-US",
- password=password2,
- )
- # Create first owner
- TenantService.create_tenant_member(tenant, account1, role="owner")
- # Try to create second owner (should fail)
- with pytest.raises(Exception, match="Tenant already has an owner"):
- TenantService.create_tenant_member(tenant, account2, role="owner")
- def test_create_tenant_member_existing_member(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test updating role for existing tenant member.
- """
- fake = Faker()
- tenant_name = fake.company()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant and account
- tenant = TenantService.create_tenant(name=tenant_name)
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Create member with initial role
- tenant_member1 = TenantService.create_tenant_member(tenant, account, role="normal")
- assert tenant_member1.role == "normal"
- # Update member role
- tenant_member2 = TenantService.create_tenant_member(tenant, account, role="editor")
- assert tenant_member2.tenant_id == tenant_member1.tenant_id
- assert tenant_member2.account_id == tenant_member1.account_id
- assert tenant_member2.role == "editor"
- def test_get_join_tenants_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test getting join tenants for an account.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- tenant1_name = fake.company()
- tenant2_name = fake.company()
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create account and tenants
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- tenant1 = TenantService.create_tenant(name=tenant1_name)
- tenant2 = TenantService.create_tenant(name=tenant2_name)
- # Add account to both tenants
- TenantService.create_tenant_member(tenant1, account, role="normal")
- TenantService.create_tenant_member(tenant2, account, role="admin")
- # Get join tenants
- join_tenants = TenantService.get_join_tenants(account)
- assert len(join_tenants) == 2
- tenant_names = [tenant.name for tenant in join_tenants]
- assert tenant1_name in tenant_names
- assert tenant2_name in tenant_names
- def test_get_current_tenant_by_account_success(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test getting current tenant by account successfully.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- tenant_name = fake.company()
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create account and tenant
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- tenant = TenantService.create_tenant(name=tenant_name)
- # Add account to tenant and set as current
- TenantService.create_tenant_member(tenant, account, role="owner")
- account.current_tenant = tenant
- from extensions.ext_database import db
- db.session.commit()
- # Get current tenant
- current_tenant = TenantService.get_current_tenant_by_account(account)
- assert current_tenant.id == tenant.id
- assert current_tenant.name == tenant.name
- assert current_tenant.role == "owner"
- def test_get_current_tenant_by_account_not_found(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test getting current tenant when account has no current tenant.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create account without setting current tenant
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Try to get current tenant (should fail)
- with pytest.raises(AttributeError):
- TenantService.get_current_tenant_by_account(account)
- def test_switch_tenant_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test successful tenant switching.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- tenant1_name = fake.company()
- tenant2_name = fake.company()
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create account and tenants
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- tenant1 = TenantService.create_tenant(name=tenant1_name)
- tenant2 = TenantService.create_tenant(name=tenant2_name)
- # Add account to both tenants
- TenantService.create_tenant_member(tenant1, account, role="owner")
- TenantService.create_tenant_member(tenant2, account, role="admin")
- # Set initial current tenant
- account.current_tenant = tenant1
- from extensions.ext_database import db
- db.session.commit()
- # Switch to second tenant
- TenantService.switch_tenant(account, tenant2.id)
- # Verify tenant was switched
- db.session.refresh(account)
- assert account.current_tenant_id == tenant2.id
- def test_switch_tenant_no_tenant_id(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test tenant switching without providing tenant ID.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Try to switch tenant without providing tenant ID
- with pytest.raises(ValueError, match="Tenant ID must be provided"):
- TenantService.switch_tenant(account, None)
- def test_switch_tenant_account_not_member(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test switching to a tenant where account is not a member.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- tenant_name = fake.company()
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create account and tenant
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- tenant = TenantService.create_tenant(name=tenant_name)
- # Try to switch to tenant where account is not a member
- with pytest.raises(Exception, match="Tenant not found or account is not a member of the tenant"):
- TenantService.switch_tenant(account, tenant.id)
- def test_has_roles_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test checking if tenant has specific roles.
- """
- fake = Faker()
- tenant_name = fake.company()
- owner_email = fake.email()
- owner_name = fake.name()
- owner_password = fake.password(length=12)
- admin_email = fake.email()
- admin_name = fake.name()
- admin_password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant and accounts
- tenant = TenantService.create_tenant(name=tenant_name)
- owner_account = AccountService.create_account(
- email=owner_email,
- name=owner_name,
- interface_language="en-US",
- password=owner_password,
- )
- admin_account = AccountService.create_account(
- email=admin_email,
- name=admin_name,
- interface_language="en-US",
- password=admin_password,
- )
- # Add members with different roles
- TenantService.create_tenant_member(tenant, owner_account, role="owner")
- TenantService.create_tenant_member(tenant, admin_account, role="admin")
- # Check if tenant has owner role
- from models.account import TenantAccountRole
- has_owner = TenantService.has_roles(tenant, [TenantAccountRole.OWNER])
- assert has_owner is True
- # Check if tenant has admin role
- has_admin = TenantService.has_roles(tenant, [TenantAccountRole.ADMIN])
- assert has_admin is True
- # Check if tenant has normal role (should be False)
- has_normal = TenantService.has_roles(tenant, [TenantAccountRole.NORMAL])
- assert has_normal is False
- def test_has_roles_invalid_role_type(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test checking roles with invalid role type.
- """
- fake = Faker()
- tenant_name = fake.company()
- invalid_role = fake.word()
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant
- tenant = TenantService.create_tenant(name=tenant_name)
- # Try to check roles with invalid role type
- with pytest.raises(ValueError, match="all roles must be TenantAccountRole"):
- TenantService.has_roles(tenant, [invalid_role])
- def test_get_user_role_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test getting user role in a tenant.
- """
- fake = Faker()
- tenant_name = fake.company()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant and account
- tenant = TenantService.create_tenant(name=tenant_name)
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Add account to tenant with specific role
- TenantService.create_tenant_member(tenant, account, role="editor")
- # Get user role
- user_role = TenantService.get_user_role(account, tenant)
- assert user_role == "editor"
- def test_check_member_permission_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test checking member permission successfully.
- """
- fake = Faker()
- tenant_name = fake.company()
- owner_email = fake.email()
- owner_name = fake.name()
- owner_password = fake.password(length=12)
- member_email = fake.email()
- member_name = fake.name()
- member_password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant and accounts
- tenant = TenantService.create_tenant(name=tenant_name)
- owner_account = AccountService.create_account(
- email=owner_email,
- name=owner_name,
- interface_language="en-US",
- password=owner_password,
- )
- member_account = AccountService.create_account(
- email=member_email,
- name=member_name,
- interface_language="en-US",
- password=member_password,
- )
- # Add members with different roles
- TenantService.create_tenant_member(tenant, owner_account, role="owner")
- TenantService.create_tenant_member(tenant, member_account, role="normal")
- # Check owner permission to add member (should succeed)
- TenantService.check_member_permission(tenant, owner_account, member_account, "add")
- def test_check_member_permission_invalid_action(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test checking member permission with invalid action.
- """
- fake = Faker()
- tenant_name = fake.company()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- invalid_action = "invalid_action_that_doesnt_exist"
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant and account
- tenant = TenantService.create_tenant(name=tenant_name)
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Add account to tenant
- TenantService.create_tenant_member(tenant, account, role="owner")
- # Try to check permission with invalid action
- with pytest.raises(Exception, match="Invalid action"):
- TenantService.check_member_permission(tenant, account, None, invalid_action)
- def test_check_member_permission_operate_self(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test checking member permission when trying to operate self.
- """
- fake = Faker()
- tenant_name = fake.company()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant and account
- tenant = TenantService.create_tenant(name=tenant_name)
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Add account to tenant
- TenantService.create_tenant_member(tenant, account, role="owner")
- # Try to check permission to operate self
- with pytest.raises(Exception, match="Cannot operate self"):
- TenantService.check_member_permission(tenant, account, account, "remove")
- def test_remove_member_from_tenant_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test successful member removal from tenant.
- """
- fake = Faker()
- tenant_name = fake.company()
- owner_email = fake.email()
- owner_name = fake.name()
- owner_password = fake.password(length=12)
- member_email = fake.email()
- member_name = fake.name()
- member_password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant and accounts
- tenant = TenantService.create_tenant(name=tenant_name)
- owner_account = AccountService.create_account(
- email=owner_email,
- name=owner_name,
- interface_language="en-US",
- password=owner_password,
- )
- member_account = AccountService.create_account(
- email=member_email,
- name=member_name,
- interface_language="en-US",
- password=member_password,
- )
- # Add members with different roles
- TenantService.create_tenant_member(tenant, owner_account, role="owner")
- TenantService.create_tenant_member(tenant, member_account, role="normal")
- # Remove member
- TenantService.remove_member_from_tenant(tenant, member_account, owner_account)
- # Verify member was removed
- from extensions.ext_database import db
- from models.account import TenantAccountJoin
- member_join = (
- db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=member_account.id).first()
- )
- assert member_join is None
- def test_remove_member_from_tenant_operate_self(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test removing member when trying to operate self.
- """
- fake = Faker()
- tenant_name = fake.company()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant and account
- tenant = TenantService.create_tenant(name=tenant_name)
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Add account to tenant
- TenantService.create_tenant_member(tenant, account, role="owner")
- # Try to remove self
- with pytest.raises(Exception, match="Cannot operate self"):
- TenantService.remove_member_from_tenant(tenant, account, account)
- def test_remove_member_from_tenant_not_member(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test removing member who is not in the tenant.
- """
- fake = Faker()
- tenant_name = fake.company()
- owner_email = fake.email()
- owner_name = fake.name()
- owner_password = fake.password(length=12)
- non_member_email = fake.email()
- non_member_name = fake.name()
- non_member_password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant and accounts
- tenant = TenantService.create_tenant(name=tenant_name)
- owner_account = AccountService.create_account(
- email=owner_email,
- name=owner_name,
- interface_language="en-US",
- password=owner_password,
- )
- non_member_account = AccountService.create_account(
- email=non_member_email,
- name=non_member_name,
- interface_language="en-US",
- password=non_member_password,
- )
- # Add only owner to tenant
- TenantService.create_tenant_member(tenant, owner_account, role="owner")
- # Try to remove non-member
- with pytest.raises(Exception, match="Member not in tenant"):
- TenantService.remove_member_from_tenant(tenant, non_member_account, owner_account)
- def test_update_member_role_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test successful member role update.
- """
- fake = Faker()
- tenant_name = fake.company()
- owner_email = fake.email()
- owner_name = fake.name()
- owner_password = fake.password(length=12)
- member_email = fake.email()
- member_name = fake.name()
- member_password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant and accounts
- tenant = TenantService.create_tenant(name=tenant_name)
- owner_account = AccountService.create_account(
- email=owner_email,
- name=owner_name,
- interface_language="en-US",
- password=owner_password,
- )
- member_account = AccountService.create_account(
- email=member_email,
- name=member_name,
- interface_language="en-US",
- password=member_password,
- )
- # Add members with different roles
- TenantService.create_tenant_member(tenant, owner_account, role="owner")
- TenantService.create_tenant_member(tenant, member_account, role="normal")
- # Update member role
- TenantService.update_member_role(tenant, member_account, "admin", owner_account)
- # Verify role was updated
- from extensions.ext_database import db
- from models.account import TenantAccountJoin
- member_join = (
- db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=member_account.id).first()
- )
- assert member_join.role == "admin"
- def test_update_member_role_to_owner(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test updating member role to owner (should change current owner to admin).
- """
- fake = Faker()
- tenant_name = fake.company()
- owner_email = fake.email()
- owner_name = fake.name()
- owner_password = fake.password(length=12)
- member_email = fake.email()
- member_name = fake.name()
- member_password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant and accounts
- tenant = TenantService.create_tenant(name=tenant_name)
- owner_account = AccountService.create_account(
- email=owner_email,
- name=owner_name,
- interface_language="en-US",
- password=owner_password,
- )
- member_account = AccountService.create_account(
- email=member_email,
- name=member_name,
- interface_language="en-US",
- password=member_password,
- )
- # Add members with different roles
- TenantService.create_tenant_member(tenant, owner_account, role="owner")
- TenantService.create_tenant_member(tenant, member_account, role="admin")
- # Update member role to owner
- TenantService.update_member_role(tenant, member_account, "owner", owner_account)
- # Verify roles were updated correctly
- from extensions.ext_database import db
- from models.account import TenantAccountJoin
- owner_join = (
- db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=owner_account.id).first()
- )
- member_join = (
- db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=member_account.id).first()
- )
- assert owner_join.role == "admin"
- assert member_join.role == "owner"
- def test_update_member_role_already_assigned(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test updating member role to already assigned role.
- """
- fake = Faker()
- tenant_name = fake.company()
- owner_email = fake.email()
- owner_name = fake.name()
- owner_password = fake.password(length=12)
- member_email = fake.email()
- member_name = fake.name()
- member_password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant and accounts
- tenant = TenantService.create_tenant(name=tenant_name)
- owner_account = AccountService.create_account(
- email=owner_email,
- name=owner_name,
- interface_language="en-US",
- password=owner_password,
- )
- member_account = AccountService.create_account(
- email=member_email,
- name=member_name,
- interface_language="en-US",
- password=member_password,
- )
- # Add members with different roles
- TenantService.create_tenant_member(tenant, owner_account, role="owner")
- TenantService.create_tenant_member(tenant, member_account, role="admin")
- # Try to update member role to already assigned role
- with pytest.raises(Exception, match="The provided role is already assigned to the member"):
- TenantService.update_member_role(tenant, member_account, "admin", owner_account)
- def test_get_tenant_count_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test getting tenant count successfully.
- """
- fake = Faker()
- tenant1_name = fake.company()
- tenant2_name = fake.company()
- tenant3_name = fake.company()
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create multiple tenants
- tenant1 = TenantService.create_tenant(name=tenant1_name)
- tenant2 = TenantService.create_tenant(name=tenant2_name)
- tenant3 = TenantService.create_tenant(name=tenant3_name)
- # Get tenant count
- tenant_count = TenantService.get_tenant_count()
- # Should have at least 3 tenants (may be more from other tests)
- assert tenant_count >= 3
- def test_create_owner_tenant_if_not_exist_new_user(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test creating owner tenant for new user without existing tenants.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- workspace_name = fake.company()
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.license.workspaces.is_available.return_value = True
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Create owner tenant
- TenantService.create_owner_tenant_if_not_exist(account, name=workspace_name)
- # Verify tenant was created and linked
- from extensions.ext_database import db
- from models.account import TenantAccountJoin
- tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first()
- assert tenant_join is not None
- assert tenant_join.role == "owner"
- assert account.current_tenant is not None
- assert account.current_tenant.name == workspace_name
- def test_create_owner_tenant_if_not_exist_existing_tenant(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test creating owner tenant when user already has a tenant.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- existing_tenant_name = fake.company()
- new_workspace_name = fake.company()
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.license.workspaces.is_available.return_value = True
- # Create account and existing tenant
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- existing_tenant = TenantService.create_tenant(name=existing_tenant_name)
- TenantService.create_tenant_member(existing_tenant, account, role="owner")
- account.current_tenant = existing_tenant
- from extensions.ext_database import db
- db.session.commit()
- # Try to create owner tenant again (should not create new one)
- TenantService.create_owner_tenant_if_not_exist(account, name=new_workspace_name)
- # Verify no new tenant was created
- tenant_joins = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).all()
- assert len(tenant_joins) == 1
- assert account.current_tenant.id == existing_tenant.id
- def test_create_owner_tenant_if_not_exist_workspace_disabled(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test creating owner tenant when workspace creation is disabled.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- workspace_name = fake.company()
- # Setup mocks to disable workspace creation
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Try to create owner tenant (should fail)
- with pytest.raises(WorkSpaceNotAllowedCreateError): # WorkSpaceNotAllowedCreateError exception
- TenantService.create_owner_tenant_if_not_exist(account, name=workspace_name)
- def test_get_tenant_members_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test getting tenant members successfully.
- """
- fake = Faker()
- tenant_name = fake.company()
- owner_email = fake.email()
- owner_name = fake.name()
- owner_password = fake.password(length=12)
- admin_email = fake.email()
- admin_name = fake.name()
- admin_password = fake.password(length=12)
- normal_email = fake.email()
- normal_name = fake.name()
- normal_password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant and accounts
- tenant = TenantService.create_tenant(name=tenant_name)
- owner_account = AccountService.create_account(
- email=owner_email,
- name=owner_name,
- interface_language="en-US",
- password=owner_password,
- )
- admin_account = AccountService.create_account(
- email=admin_email,
- name=admin_name,
- interface_language="en-US",
- password=admin_password,
- )
- normal_account = AccountService.create_account(
- email=normal_email,
- name=normal_name,
- interface_language="en-US",
- password=normal_password,
- )
- # Add members with different roles
- TenantService.create_tenant_member(tenant, owner_account, role="owner")
- TenantService.create_tenant_member(tenant, admin_account, role="admin")
- TenantService.create_tenant_member(tenant, normal_account, role="normal")
- # Get tenant members
- members = TenantService.get_tenant_members(tenant)
- assert len(members) == 3
- member_emails = [member.email for member in members]
- assert owner_email in member_emails
- assert admin_email in member_emails
- assert normal_email in member_emails
- # Verify roles are set correctly
- for member in members:
- if member.email == owner_email:
- assert member.role == "owner"
- elif member.email == admin_email:
- assert member.role == "admin"
- elif member.email == normal_email:
- assert member.role == "normal"
- def test_get_dataset_operator_members_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test getting dataset operator members successfully.
- """
- fake = Faker()
- tenant_name = fake.company()
- owner_email = fake.email()
- owner_name = fake.name()
- owner_password = fake.password(length=12)
- operator_email = fake.email()
- operator_name = fake.name()
- operator_password = fake.password(length=12)
- normal_email = fake.email()
- normal_name = fake.name()
- normal_password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant and accounts
- tenant = TenantService.create_tenant(name=tenant_name)
- owner_account = AccountService.create_account(
- email=owner_email,
- name=owner_name,
- interface_language="en-US",
- password=owner_password,
- )
- dataset_operator_account = AccountService.create_account(
- email=operator_email,
- name=operator_name,
- interface_language="en-US",
- password=operator_password,
- )
- normal_account = AccountService.create_account(
- email=normal_email,
- name=normal_name,
- interface_language="en-US",
- password=normal_password,
- )
- # Add members with different roles
- TenantService.create_tenant_member(tenant, owner_account, role="owner")
- TenantService.create_tenant_member(tenant, dataset_operator_account, role="dataset_operator")
- TenantService.create_tenant_member(tenant, normal_account, role="normal")
- # Get dataset operator members
- dataset_operators = TenantService.get_dataset_operator_members(tenant)
- assert len(dataset_operators) == 1
- assert dataset_operators[0].email == operator_email
- assert dataset_operators[0].role == "dataset_operator"
- def test_get_custom_config_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test getting custom config successfully.
- """
- fake = Faker()
- tenant_name = fake.company()
- theme = fake.random_element(elements=("dark", "light"))
- language = fake.random_element(elements=("zh-CN", "en-US"))
- # Setup mocks
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- # Create tenant with custom config
- tenant = TenantService.create_tenant(name=tenant_name)
- # Set custom config
- custom_config = {"theme": theme, "language": language, "feature_flags": {"beta": True}}
- tenant.custom_config_dict = custom_config
- from extensions.ext_database import db
- db.session.commit()
- # Get custom config
- retrieved_config = TenantService.get_custom_config(tenant.id)
- assert retrieved_config == custom_config
- assert retrieved_config["theme"] == theme
- assert retrieved_config["language"] == language
- assert retrieved_config["feature_flags"]["beta"] is True
- class TestRegisterService:
- """Integration tests for RegisterService using testcontainers."""
- @pytest.fixture
- def mock_external_service_dependencies(self):
- """Mock setup for external service dependencies."""
- with (
- patch("services.account_service.FeatureService") as mock_feature_service,
- patch("services.account_service.BillingService") as mock_billing_service,
- patch("services.account_service.PassportService") as mock_passport_service,
- ):
- # Setup default mock returns
- mock_feature_service.get_system_features.return_value.is_allow_register = True
- mock_feature_service.get_system_features.return_value.is_allow_create_workspace = True
- mock_feature_service.get_system_features.return_value.license.workspaces.is_available.return_value = True
- mock_billing_service.is_email_in_freeze.return_value = False
- mock_passport_service.return_value.issue.return_value = "mock_jwt_token"
- yield {
- "feature_service": mock_feature_service,
- "billing_service": mock_billing_service,
- "passport_service": mock_passport_service,
- }
- def test_setup_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test successful system setup with account creation and tenant setup.
- """
- fake = Faker()
- admin_email = fake.email()
- admin_name = fake.name()
- admin_password = fake.password(length=12)
- ip_address = fake.ipv4()
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Execute setup
- RegisterService.setup(
- email=admin_email,
- name=admin_name,
- password=admin_password,
- ip_address=ip_address,
- )
- # Verify account was created
- from extensions.ext_database import db
- from models.account import Account
- from models.model import DifySetup
- account = db.session.query(Account).filter_by(email=admin_email).first()
- assert account is not None
- assert account.name == admin_name
- assert account.last_login_ip == ip_address
- assert account.initialized_at is not None
- assert account.status == "active"
- # Verify DifySetup was created
- dify_setup = db.session.query(DifySetup).first()
- assert dify_setup is not None
- # Verify tenant was created and linked
- from models.account import TenantAccountJoin
- tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first()
- assert tenant_join is not None
- assert tenant_join.role == "owner"
- def test_setup_failure_rollback(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test setup failure with proper rollback of all created entities.
- """
- fake = Faker()
- admin_email = fake.email()
- admin_name = fake.name()
- admin_password = fake.password(length=12)
- ip_address = fake.ipv4()
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Mock AccountService.create_account to raise exception
- with patch("services.account_service.AccountService.create_account") as mock_create_account:
- mock_create_account.side_effect = Exception("Database error")
- # Execute setup and verify exception
- with pytest.raises(ValueError, match="Setup failed: Database error"):
- RegisterService.setup(
- email=admin_email,
- name=admin_name,
- password=admin_password,
- ip_address=ip_address,
- )
- # Verify no entities were created (rollback worked)
- from extensions.ext_database import db
- from models.account import Account, Tenant, TenantAccountJoin
- from models.model import DifySetup
- account = db.session.query(Account).filter_by(email=admin_email).first()
- tenant_count = db.session.query(Tenant).count()
- tenant_join_count = db.session.query(TenantAccountJoin).count()
- dify_setup_count = db.session.query(DifySetup).count()
- assert account is None
- assert tenant_count == 0
- assert tenant_join_count == 0
- assert dify_setup_count == 0
- def test_register_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test successful account registration with workspace creation.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- language = fake.random_element(elements=("en-US", "zh-CN"))
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.license.workspaces.is_available.return_value = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Execute registration
- account = RegisterService.register(
- email=email,
- name=name,
- password=password,
- language=language,
- )
- # Verify account was created
- assert account.email == email
- assert account.name == name
- assert account.status == "active"
- assert account.initialized_at is not None
- # Verify tenant was created and linked
- from extensions.ext_database import db
- from models.account import TenantAccountJoin
- tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first()
- assert tenant_join is not None
- assert tenant_join.role == "owner"
- assert account.current_tenant is not None
- assert account.current_tenant.name == f"{name}'s Workspace"
- def test_register_with_oauth(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test account registration with OAuth integration.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- open_id = fake.uuid4()
- provider = fake.random_element(elements=("google", "github", "microsoft"))
- language = fake.random_element(elements=("en-US", "zh-CN"))
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.license.workspaces.is_available.return_value = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Execute registration with OAuth
- account = RegisterService.register(
- email=email,
- name=name,
- password=None,
- open_id=open_id,
- provider=provider,
- language=language,
- )
- # Verify account was created
- assert account.email == email
- assert account.name == name
- assert account.status == "active"
- assert account.initialized_at is not None
- # Verify OAuth integration was created
- from extensions.ext_database import db
- from models.account import AccountIntegrate
- integration = db.session.query(AccountIntegrate).filter_by(account_id=account.id, provider=provider).first()
- assert integration is not None
- assert integration.open_id == open_id
- def test_register_with_pending_status(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test account registration with pending status.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- language = fake.random_element(elements=("en-US", "zh-CN"))
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.license.workspaces.is_available.return_value = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Execute registration with pending status
- from models.account import AccountStatus
- account = RegisterService.register(
- email=email,
- name=name,
- password=password,
- language=language,
- status=AccountStatus.PENDING,
- )
- # Verify account was created with pending status
- assert account.email == email
- assert account.name == name
- assert account.status == "pending"
- assert account.initialized_at is not None
- # Verify tenant was created and linked
- from extensions.ext_database import db
- from models.account import TenantAccountJoin
- tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first()
- assert tenant_join is not None
- assert tenant_join.role == "owner"
- def test_register_workspace_creation_disabled(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test account registration when workspace creation is disabled.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- language = fake.random_element(elements=("en-US", "zh-CN"))
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = False
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # with pytest.raises(AccountRegisterError, match="Workspace is not allowed to create."):
- account = RegisterService.register(
- email=email,
- name=name,
- password=password,
- language=language,
- )
- # Verify account was created with no tenant
- assert account.email == email
- assert account.name == name
- assert account.status == "active"
- assert account.initialized_at is not None
- # Verify tenant was created and linked
- from extensions.ext_database import db
- from models.account import TenantAccountJoin
- tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first()
- assert tenant_join is None
- def test_register_workspace_limit_exceeded(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test account registration when workspace limit is exceeded.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- language = fake.random_element(elements=("en-US", "zh-CN"))
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.license.workspaces.is_available.return_value = False
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # with pytest.raises(AccountRegisterError, match="Workspace is not allowed to create."):
- account = RegisterService.register(
- email=email,
- name=name,
- password=password,
- language=language,
- )
- # Verify account was created with no tenant
- assert account.email == email
- assert account.name == name
- assert account.status == "active"
- assert account.initialized_at is not None
- # Verify tenant was created and linked
- from extensions.ext_database import db
- from models.account import TenantAccountJoin
- tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first()
- assert tenant_join is None
- def test_register_without_workspace(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test account registration without workspace creation.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- language = fake.random_element(elements=("en-US", "zh-CN"))
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Execute registration without workspace creation
- account = RegisterService.register(
- email=email,
- name=name,
- password=password,
- language=language,
- create_workspace_required=False,
- )
- # Verify account was created
- assert account.email == email
- assert account.name == name
- assert account.status == "active"
- assert account.initialized_at is not None
- # Verify no tenant was created
- from extensions.ext_database import db
- from models.account import TenantAccountJoin
- tenant_join = db.session.query(TenantAccountJoin).filter_by(account_id=account.id).first()
- assert tenant_join is None
- def test_invite_new_member_new_account(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test inviting a new member who doesn't have an account yet.
- """
- fake = Faker()
- tenant_name = fake.company()
- inviter_email = fake.email()
- inviter_name = fake.name()
- inviter_password = fake.password(length=12)
- new_member_email = fake.email()
- language = fake.random_element(elements=("en-US", "zh-CN"))
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.is_allow_create_workspace = True
- mock_external_service_dependencies[
- "feature_service"
- ].get_system_features.return_value.license.workspaces.is_available.return_value = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create tenant and inviter account
- tenant = TenantService.create_tenant(name=tenant_name)
- inviter = AccountService.create_account(
- email=inviter_email,
- name=inviter_name,
- interface_language="en-US",
- password=inviter_password,
- )
- TenantService.create_tenant_member(tenant, inviter, role="owner")
- # Mock the email task
- with patch("services.account_service.send_invite_member_mail_task") as mock_send_mail:
- mock_send_mail.delay.return_value = None
- # Execute invitation
- token = RegisterService.invite_new_member(
- tenant=tenant,
- email=new_member_email,
- language=language,
- role="normal",
- inviter=inviter,
- )
- # Verify token was generated
- assert token is not None
- assert len(token) > 0
- # Verify email task was called
- mock_send_mail.delay.assert_called_once()
- # Verify new account was created with pending status
- from extensions.ext_database import db
- from models.account import Account, TenantAccountJoin
- new_account = db.session.query(Account).filter_by(email=new_member_email).first()
- assert new_account is not None
- assert new_account.name == new_member_email.split("@")[0] # Default name from email
- assert new_account.status == "pending"
- # Verify tenant member was created
- tenant_join = (
- db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=new_account.id).first()
- )
- assert tenant_join is not None
- assert tenant_join.role == "normal"
- def test_invite_new_member_existing_account(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test inviting an existing member who is not in the tenant yet.
- """
- fake = Faker()
- tenant_name = fake.company()
- inviter_email = fake.email()
- inviter_name = fake.name()
- inviter_password = fake.password(length=12)
- existing_member_email = fake.email()
- existing_member_name = fake.name()
- existing_member_password = fake.password(length=12)
- language = fake.random_element(elements=("en-US", "zh-CN"))
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create tenant and inviter account
- tenant = TenantService.create_tenant(name=tenant_name)
- inviter = AccountService.create_account(
- email=inviter_email,
- name=inviter_name,
- interface_language="en-US",
- password=inviter_password,
- )
- TenantService.create_tenant_member(tenant, inviter, role="owner")
- # Create existing account
- existing_account = AccountService.create_account(
- email=existing_member_email,
- name=existing_member_name,
- interface_language="en-US",
- password=existing_member_password,
- )
- # Mock the email task
- with patch("services.account_service.send_invite_member_mail_task") as mock_send_mail:
- mock_send_mail.delay.return_value = None
- with pytest.raises(AccountAlreadyInTenantError, match="Account already in tenant."):
- # Execute invitation
- token = RegisterService.invite_new_member(
- tenant=tenant,
- email=existing_member_email,
- language=language,
- role="admin",
- inviter=inviter,
- )
- # Verify email task was not called
- mock_send_mail.delay.assert_not_called()
- # Verify tenant member was created for existing account
- from extensions.ext_database import db
- from models.account import TenantAccountJoin
- tenant_join = (
- db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=existing_account.id).first()
- )
- assert tenant_join is not None
- assert tenant_join.role == "admin"
- def test_invite_new_member_existing_member(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test inviting a member who is already in the tenant with pending status.
- """
- fake = Faker()
- tenant_name = fake.company()
- inviter_email = fake.email()
- inviter_name = fake.name()
- inviter_password = fake.password(length=12)
- existing_pending_member_email = fake.email()
- existing_pending_member_name = fake.name()
- existing_pending_member_password = fake.password(length=12)
- language = fake.random_element(elements=("en-US", "zh-CN"))
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create tenant and inviter account
- tenant = TenantService.create_tenant(name=tenant_name)
- inviter = AccountService.create_account(
- email=inviter_email,
- name=inviter_name,
- interface_language="en-US",
- password=inviter_password,
- )
- TenantService.create_tenant_member(tenant, inviter, role="owner")
- # Create existing account with pending status
- existing_account = AccountService.create_account(
- email=existing_pending_member_email,
- name=existing_pending_member_name,
- interface_language="en-US",
- password=existing_pending_member_password,
- )
- existing_account.status = "pending"
- from extensions.ext_database import db
- db.session.commit()
- # Add existing account to tenant
- TenantService.create_tenant_member(tenant, existing_account, role="normal")
- # Mock the email task
- with patch("services.account_service.send_invite_member_mail_task") as mock_send_mail:
- mock_send_mail.delay.return_value = None
- # Execute invitation (should resend email for pending member)
- token = RegisterService.invite_new_member(
- tenant=tenant,
- email=existing_pending_member_email,
- language=language,
- role="normal",
- inviter=inviter,
- )
- # Verify token was generated
- assert token is not None
- assert len(token) > 0
- # Verify email task was called
- mock_send_mail.delay.assert_called_once()
- def test_invite_new_member_no_inviter(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test inviting a member without providing an inviter.
- """
- fake = Faker()
- tenant_name = fake.company()
- new_member_email = fake.email()
- language = fake.random_element(elements=("en-US", "zh-CN"))
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create tenant
- tenant = TenantService.create_tenant(name=tenant_name)
- # Execute invitation without inviter (should fail)
- with pytest.raises(ValueError, match="Inviter is required"):
- RegisterService.invite_new_member(
- tenant=tenant,
- email=new_member_email,
- language=language,
- role="normal",
- inviter=None,
- )
- def test_invite_new_member_account_already_in_tenant(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test inviting a member who is already in the tenant with active status.
- """
- fake = Faker()
- tenant_name = fake.company()
- inviter_email = fake.email()
- inviter_name = fake.name()
- inviter_password = fake.password(length=12)
- already_in_tenant_email = fake.email()
- already_in_tenant_name = fake.name()
- already_in_tenant_password = fake.password(length=12)
- language = fake.random_element(elements=("en-US", "zh-CN"))
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create tenant and inviter account
- tenant = TenantService.create_tenant(name=tenant_name)
- inviter = AccountService.create_account(
- email=inviter_email,
- name=inviter_name,
- interface_language="en-US",
- password=inviter_password,
- )
- TenantService.create_tenant_member(tenant, inviter, role="owner")
- # Create existing account with active status
- existing_account = AccountService.create_account(
- email=already_in_tenant_email,
- name=already_in_tenant_name,
- interface_language="en-US",
- password=already_in_tenant_password,
- )
- existing_account.status = "active"
- from extensions.ext_database import db
- db.session.commit()
- # Add existing account to tenant
- TenantService.create_tenant_member(tenant, existing_account, role="normal")
- # Execute invitation (should fail for active member)
- with pytest.raises(AccountAlreadyInTenantError, match="Account already in tenant."):
- RegisterService.invite_new_member(
- tenant=tenant,
- email=already_in_tenant_email,
- language=language,
- role="normal",
- inviter=inviter,
- )
- def test_generate_invite_token_success(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test successful generation of invite token.
- """
- fake = Faker()
- tenant_name = fake.company()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create tenant and account
- tenant = TenantService.create_tenant(name=tenant_name)
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Execute token generation
- token = RegisterService.generate_invite_token(tenant, account)
- # Verify token was generated
- assert token is not None
- assert len(token) > 0
- # Verify token was stored in Redis
- from extensions.ext_redis import redis_client
- token_key = RegisterService._get_invitation_token_key(token)
- stored_data = redis_client.get(token_key)
- assert stored_data is not None
- # Verify stored data contains correct information
- import json
- invitation_data = json.loads(stored_data.decode("utf-8"))
- assert invitation_data["account_id"] == str(account.id)
- assert invitation_data["email"] == account.email
- assert invitation_data["workspace_id"] == tenant.id
- def test_is_valid_invite_token_valid(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test validation of valid invite token.
- """
- fake = Faker()
- tenant_name = fake.company()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create tenant and account
- tenant = TenantService.create_tenant(name=tenant_name)
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Generate a real token
- token = RegisterService.generate_invite_token(tenant, account)
- # Execute validation
- is_valid = RegisterService.is_valid_invite_token(token)
- # Verify token is valid
- assert is_valid is True
- def test_is_valid_invite_token_invalid(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test validation of invalid invite token.
- """
- fake = Faker()
- invalid_token = fake.uuid4()
- # Execute validation with non-existent token
- is_valid = RegisterService.is_valid_invite_token(invalid_token)
- # Verify token is invalid
- assert is_valid is False
- def test_revoke_token_with_workspace_and_email(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test revoking token with workspace ID and email.
- """
- fake = Faker()
- tenant_name = fake.company()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create tenant and account
- tenant = TenantService.create_tenant(name=tenant_name)
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Generate a real token
- token = RegisterService.generate_invite_token(tenant, account)
- # Verify token exists in Redis before revocation
- from extensions.ext_redis import redis_client
- token_key = RegisterService._get_invitation_token_key(token)
- assert redis_client.get(token_key) is not None
- # Execute token revocation
- RegisterService.revoke_token(
- workspace_id=tenant.id,
- email=account.email,
- token=token,
- )
- # Verify token was not deleted from Redis
- assert redis_client.get(token_key) is not None
- def test_revoke_token_without_workspace_and_email(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test revoking token without workspace ID and email.
- """
- fake = Faker()
- tenant_name = fake.company()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create tenant and account
- tenant = TenantService.create_tenant(name=tenant_name)
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Generate a real token
- token = RegisterService.generate_invite_token(tenant, account)
- # Verify token exists in Redis before revocation
- from extensions.ext_redis import redis_client
- token_key = RegisterService._get_invitation_token_key(token)
- assert redis_client.get(token_key) is not None
- # Execute token revocation without workspace and email
- RegisterService.revoke_token(
- workspace_id="",
- email="",
- token=token,
- )
- # Verify token was deleted from Redis
- assert redis_client.get(token_key) is None
- def test_get_invitation_if_token_valid_success(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test getting invitation data with valid token.
- """
- fake = Faker()
- tenant_name = fake.company()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create tenant and account
- tenant = TenantService.create_tenant(name=tenant_name)
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- TenantService.create_tenant_member(tenant, account, role="normal")
- # Generate a real token
- token = RegisterService.generate_invite_token(tenant, account)
- email_hash = sha256(account.email.encode()).hexdigest()
- cache_key = f"member_invite_token:{tenant.id}, {email_hash}:{token}"
- from extensions.ext_redis import redis_client
- redis_client.setex(cache_key, 24 * 60 * 60, account.id)
- # Execute invitation retrieval
- result = RegisterService.get_invitation_if_token_valid(
- workspace_id=tenant.id,
- email=account.email,
- token=token,
- )
- # Verify result contains expected data
- assert result is not None
- assert result["account"].id == account.id
- assert result["tenant"].id == tenant.id
- assert result["data"]["account_id"] == str(account.id)
- assert result["data"]["email"] == account.email
- assert result["data"]["workspace_id"] == tenant.id
- def test_get_invitation_if_token_valid_invalid_token(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test getting invitation data with invalid token.
- """
- fake = Faker()
- workspace_id = fake.uuid4()
- email = fake.email()
- invalid_token = fake.uuid4()
- # Execute invitation retrieval with invalid token
- result = RegisterService.get_invitation_if_token_valid(
- workspace_id=workspace_id,
- email=email,
- token=invalid_token,
- )
- # Verify result is None
- assert result is None
- def test_get_invitation_if_token_valid_invalid_tenant(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test getting invitation data with invalid tenant.
- """
- fake = Faker()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- invalid_tenant_id = fake.uuid4()
- token = fake.uuid4()
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create account
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- # Create a real token but with non-existent tenant ID
- from extensions.ext_redis import redis_client
- invitation_data = {
- "account_id": str(account.id),
- "email": account.email,
- "workspace_id": invalid_tenant_id,
- }
- token_key = RegisterService._get_invitation_token_key(token)
- import json
- redis_client.setex(token_key, 24 * 60 * 60, json.dumps(invitation_data))
- # Execute invitation retrieval
- result = RegisterService.get_invitation_if_token_valid(
- workspace_id=invalid_tenant_id,
- email=account.email,
- token=token,
- )
- # Verify result is None (tenant not found)
- assert result is None
- # Clean up
- redis_client.delete(token_key)
- def test_get_invitation_if_token_valid_account_mismatch(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test getting invitation data with account ID mismatch.
- """
- fake = Faker()
- tenant_name = fake.company()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- token = fake.uuid4()
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create tenant and account
- tenant = TenantService.create_tenant(name=tenant_name)
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- TenantService.create_tenant_member(tenant, account, role="normal")
- # Create a real token but with mismatched account ID
- from extensions.ext_redis import redis_client
- invitation_data = {
- "account_id": "different-account-id", # Different from actual account ID
- "email": account.email,
- "workspace_id": tenant.id,
- }
- token_key = RegisterService._get_invitation_token_key(token)
- redis_client.setex(token_key, 24 * 60 * 60, json.dumps(invitation_data))
- # Execute invitation retrieval
- result = RegisterService.get_invitation_if_token_valid(
- workspace_id=tenant.id,
- email=account.email,
- token=token,
- )
- # Verify result is None (account ID mismatch)
- assert result is None
- # Clean up
- redis_client.delete(token_key)
- def test_get_invitation_if_token_valid_tenant_not_normal(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test getting invitation data with tenant not in normal status.
- """
- fake = Faker()
- tenant_name = fake.company()
- email = fake.email()
- name = fake.name()
- password = fake.password(length=12)
- token = fake.uuid4()
- # Setup mocks
- mock_external_service_dependencies["feature_service"].get_system_features.return_value.is_allow_register = True
- mock_external_service_dependencies["billing_service"].is_email_in_freeze.return_value = False
- # Create tenant and account
- tenant = TenantService.create_tenant(name=tenant_name)
- account = AccountService.create_account(
- email=email,
- name=name,
- interface_language="en-US",
- password=password,
- )
- TenantService.create_tenant_member(tenant, account, role="normal")
- # Change tenant status to non-normal
- tenant.status = "suspended"
- from extensions.ext_database import db
- db.session.commit()
- # Create a real token
- from extensions.ext_redis import redis_client
- invitation_data = {
- "account_id": str(account.id),
- "email": account.email,
- "workspace_id": tenant.id,
- }
- token_key = RegisterService._get_invitation_token_key(token)
- import json
- redis_client.setex(token_key, 24 * 60 * 60, json.dumps(invitation_data))
- # Execute invitation retrieval
- result = RegisterService.get_invitation_if_token_valid(
- workspace_id=tenant.id,
- email=account.email,
- token=token,
- )
- # Verify result is None (tenant not in normal status)
- assert result is None
- # Clean up
- redis_client.delete(token_key)
- def test_get_invitation_by_token_with_workspace_and_email(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test getting invitation by token with workspace ID and email.
- """
- fake = Faker()
- token = fake.uuid4()
- workspace_id = fake.uuid4()
- email = fake.email()
- # Create the cache key as the service does
- from hashlib import sha256
- from extensions.ext_redis import redis_client
- email_hash = sha256(email.encode()).hexdigest()
- cache_key = f"member_invite_token:{workspace_id}, {email_hash}:{token}"
- # Store account ID in Redis
- account_id = fake.uuid4()
- redis_client.setex(cache_key, 24 * 60 * 60, account_id)
- # Execute invitation retrieval
- result = RegisterService._get_invitation_by_token(
- token=token,
- workspace_id=workspace_id,
- email=email,
- )
- # Verify result contains expected data
- assert result is not None
- assert result["account_id"] == account_id
- assert result["email"] == email
- assert result["workspace_id"] == workspace_id
- # Clean up
- redis_client.delete(cache_key)
- def test_get_invitation_by_token_without_workspace_and_email(
- self, db_session_with_containers, mock_external_service_dependencies
- ):
- """
- Test getting invitation by token without workspace ID and email.
- """
- fake = Faker()
- token = fake.uuid4()
- invitation_data = {
- "account_id": fake.uuid4(),
- "email": fake.email(),
- "workspace_id": fake.uuid4(),
- }
- # Store invitation data in Redis using standard token key
- from extensions.ext_redis import redis_client
- token_key = RegisterService._get_invitation_token_key(token)
- import json
- redis_client.setex(token_key, 24 * 60 * 60, json.dumps(invitation_data))
- # Execute invitation retrieval
- result = RegisterService._get_invitation_by_token(token=token)
- # Verify result contains expected data
- assert result is not None
- assert result["account_id"] == invitation_data["account_id"]
- assert result["email"] == invitation_data["email"]
- assert result["workspace_id"] == invitation_data["workspace_id"]
- # Clean up
- redis_client.delete(token_key)
- def test_get_invitation_token_key(self, db_session_with_containers, mock_external_service_dependencies):
- """
- Test getting invitation token key.
- """
- fake = Faker()
- token = fake.uuid4()
- # Execute token key generation
- token_key = RegisterService._get_invitation_token_key(token)
- # Verify token key format
- assert token_key == f"member_invite:token:{token}"
|