From b9d8f11f2dbcbd9f10bbd581e229210e6d1afad0 Mon Sep 17 00:00:00 2001
From: Son NK <>
Date: Thu, 31 Dec 2020 13:59:03 +0100
Subject: [PATCH] put the default domain to top

---
 app/dashboard/views/custom_alias.py  | 41 ++++++++++++++++++++--------
 tests/dashboard/test_custom_alias.py | 22 ++++++++++++++-
 tests/utils.py                       |  7 +++--
 3 files changed, 56 insertions(+), 14 deletions(-)

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"),