From 9874422700e03fa04d068f83561b00041ccbe3ff Mon Sep 17 00:00:00 2001
From: Son NK <>
Date: Sat, 2 May 2020 12:34:11 +0200
Subject: [PATCH] refactor custom_alias: create available_suffixes()

---
 app/dashboard/views/custom_alias.py  | 41 +++++++++++++++++++---------
 tests/dashboard/test_custom_alias.py | 21 +++++++++++++-
 2 files changed, 48 insertions(+), 14 deletions(-)

diff --git a/app/dashboard/views/custom_alias.py b/app/dashboard/views/custom_alias.py
index 4523f2fb..b7fda5da 100644
--- a/app/dashboard/views/custom_alias.py
+++ b/app/dashboard/views/custom_alias.py
@@ -11,12 +11,32 @@ from app.dashboard.base import dashboard_bp
 from app.email_utils import email_belongs_to_alias_domains, get_email_domain_part
 from app.extensions import db
 from app.log import LOG
-from app.models import Alias, CustomDomain, DeletedAlias, Mailbox
+from app.models import Alias, CustomDomain, DeletedAlias, Mailbox, User
 from app.utils import convert_to_id, random_word, word_exist
 
 signer = TimestampSigner(CUSTOM_ALIAS_SECRET)
 
 
+def available_suffixes(user: User) -> [bool, str, str]:
+    """Return (is_custom_domain, alias-suffix, time-signed alias-suffix)"""
+    user_custom_domains = [cd.domain for cd in user.verified_custom_domains()]
+
+    # List of (is_custom_domain, alias-suffix, time-signed alias-suffix)
+    suffixes = []
+
+    # put custom domain first
+    for alias_domain in user_custom_domains:
+        suffix = "@" + alias_domain
+        suffixes.append((True, suffix, signer.sign(suffix).decode()))
+
+    # then default domain
+    for domain in ALIAS_DOMAINS:
+        suffix = ("" if DISABLE_ALIAS_SUFFIX else "." + random_word()) + "@" + domain
+        suffixes.append((False, suffix, signer.sign(suffix).decode()))
+
+    return suffixes
+
+
 @dashboard_bp.route("/custom_alias", methods=["GET", "POST"])
 @login_required
 def custom_alias():
@@ -32,17 +52,7 @@ def custom_alias():
 
     user_custom_domains = [cd.domain for cd in current_user.verified_custom_domains()]
     # List of (is_custom_domain, alias-suffix, time-signed alias-suffix)
-    suffixes = []
-
-    # put custom domain first
-    for alias_domain in user_custom_domains:
-        suffix = "@" + alias_domain
-        suffixes.append((True, suffix, signer.sign(suffix).decode()))
-
-    # then default domain
-    for domain in ALIAS_DOMAINS:
-        suffix = ("" if DISABLE_ALIAS_SUFFIX else "." + random_word()) + "@" + domain
-        suffixes.append((False, suffix, signer.sign(suffix).decode()))
+    suffixes = available_suffixes(current_user)
 
     mailboxes = [mb.email for mb in current_user.mailboxes()]
 
@@ -105,7 +115,12 @@ def custom_alias():
         else:
             flash("something went wrong", "warning")
 
-    return render_template("dashboard/custom_alias.html", **locals())
+    return render_template(
+        "dashboard/custom_alias.html",
+        user_custom_domains=user_custom_domains,
+        suffixes=suffixes,
+        mailboxes=mailboxes,
+    )
 
 
 def verify_prefix_suffix(user, alias_prefix, alias_suffix) -> bool:
diff --git a/tests/dashboard/test_custom_alias.py b/tests/dashboard/test_custom_alias.py
index 7245d385..3e189c3b 100644
--- a/tests/dashboard/test_custom_alias.py
+++ b/tests/dashboard/test_custom_alias.py
@@ -1,7 +1,11 @@
 from flask import url_for
 
 from app.config import EMAIL_DOMAIN
-from app.dashboard.views.custom_alias import signer, verify_prefix_suffix
+from app.dashboard.views.custom_alias import (
+    signer,
+    verify_prefix_suffix,
+    available_suffixes,
+)
 from app.extensions import db
 from app.models import Mailbox, CustomDomain
 from app.utils import random_word
@@ -53,3 +57,18 @@ def test_verify_prefix_suffix(flask_client):
     word = random_word()
     suffix = f".{word}@{EMAIL_DOMAIN}"
     assert verify_prefix_suffix(user, "prefix", suffix)
+
+
+def test_available_suffixes(flask_client):
+    user = login(flask_client)
+    db.session.commit()
+
+    CustomDomain.create(user_id=user.id, domain="test.com", verified=True)
+
+    assert len(available_suffixes(user)) > 0
+
+    # first suffix is custom domain
+    first_suffix = available_suffixes(user)[0]
+    assert first_suffix[0]
+    assert first_suffix[1] == "@test.com"
+    assert first_suffix[2].startswith("@test.com")