diff --git a/app/dashboard/views/custom_alias.py b/app/dashboard/views/custom_alias.py index cb09279d..03fffa62 100644 --- a/app/dashboard/views/custom_alias.py +++ b/app/dashboard/views/custom_alias.py @@ -39,16 +39,29 @@ def available_suffixes(user: User) -> [bool, str, str]: for alias_domain in user_custom_domains: suffix = "@" + alias_domain.domain suffix_info = (True, suffix, signer.sign(suffix).decode()) - suffixes.append(suffix_info) + + # put the default domain to top + if user.default_random_alias_domain_id == alias_domain.id: + suffixes.insert(0, suffix_info) + else: + suffixes.append(suffix_info) + if alias_domain.random_prefix_generation: suffix = "." + random_word() + "@" + alias_domain.domain suffixes.append((True, suffix, signer.sign(suffix).decode())) # then SimpleLogin domain - for domain in user.available_sl_domains(): - suffix = ("" if DISABLE_ALIAS_SUFFIX else "." + random_word()) + "@" + domain + for domain in user.get_sl_domains(): + suffix = ( + ("" if DISABLE_ALIAS_SUFFIX else "." + random_word()) + "@" + domain.domain + ) suffix_info = (False, suffix, signer.sign(suffix).decode()) - suffixes.append(suffix_info) + + # put the default domain to top + if user.default_random_alias_public_domain_id == domain.id: + suffixes.insert(0, suffix_info) + else: + suffixes.append(suffix_info) return suffixes @@ -82,9 +95,12 @@ def available_suffixes_more_info(user: User) -> [SuffixInfo]: suffixes.append(SuffixInfo(True, suffix, signer.sign(suffix).decode(), False)) if alias_domain.random_prefix_generation: suffix = "." + random_word() + "@" + alias_domain.domain - suffixes.append( - SuffixInfo(True, suffix, signer.sign(suffix).decode(), False) - ) + suffix_info = SuffixInfo(True, suffix, signer.sign(suffix).decode(), False) + # put the default domain to top + if user.default_random_alias_domain_id == alias_domain.id: + suffixes.insert(0, suffix_info) + else: + suffixes.append(suffix_info) # then SimpleLogin domain for sl_domain in user.get_sl_domains(): @@ -93,11 +109,14 @@ def available_suffixes_more_info(user: User) -> [SuffixInfo]: + "@" + sl_domain.domain ) - suffixes.append( - SuffixInfo( - False, suffix, signer.sign(suffix).decode(), sl_domain.premium_only - ) + suffix_info = SuffixInfo( + False, suffix, signer.sign(suffix).decode(), sl_domain.premium_only ) + # put the default domain to top + if user.default_random_alias_public_domain_id == sl_domain.id: + suffixes.insert(0, suffix_info) + else: + suffixes.append(suffix_info) return suffixes diff --git a/tests/dashboard/test_custom_alias.py b/tests/dashboard/test_custom_alias.py index 8a24296b..d5ae63df 100644 --- a/tests/dashboard/test_custom_alias.py +++ b/tests/dashboard/test_custom_alias.py @@ -15,6 +15,7 @@ from app.models import ( User, DomainDeletedAlias, DeletedAlias, + SLDomain, ) from app.utils import random_word from tests.utils import login @@ -104,7 +105,6 @@ def test_verify_prefix_suffix(flask_client): def test_available_suffixes(flask_client): user = login(flask_client) - db.session.commit() CustomDomain.create(user_id=user.id, domain="test.com", verified=True) @@ -117,6 +117,26 @@ def test_available_suffixes(flask_client): assert first_suffix[2].startswith("@test.com") +def test_available_suffixes_default_domain(flask_client): + user = login(flask_client) + + sl_domain = SLDomain.query.first() + CustomDomain.create( + user_id=user.id, domain="test.com", verified=True, commit=True + ) + + user.default_random_alias_public_domain_id = sl_domain.id + + # first suffix is SL Domain + first_suffix = available_suffixes(user)[0] + assert first_suffix[1].endswith(f"@{sl_domain.domain}") + + user.default_random_alias_public_domain_id = None + # first suffix is custom domain + first_suffix = available_suffixes(user)[0] + assert first_suffix[1] == "@test.com" + + def test_add_already_existed_alias(flask_client): user = login(flask_client) db.session.commit() diff --git a/tests/utils.py b/tests/utils.py index cc603627..1a086f61 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -9,9 +9,12 @@ from app.models import User def login(flask_client) -> User: # create user, user is activated user = User.create( - email="a@b.c", password="password", name="Test User", activated=True + email="a@b.c", + password="password", + name="Test User", + activated=True, + commit=True, ) - db.session.commit() r = flask_client.post( url_for("auth.login"),