From a5394286078b1676b01df6bff47aa41c94cd9a43 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Adri=C3=A0=20Casaj=C3=BAs?=
 <acasajus@users.noreply.github.com>
Date: Thu, 1 Jun 2023 10:46:25 +0200
Subject: [PATCH] Fix: If default domain is premium for free users do not offer
 it as an option (#1763)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* Fix: If default domain is premium for free users do not offer it as an option

* Refactored into simpler logic

---------

Co-authored-by: Adrià Casajús <adria.casajus@proton.ch>
---
 app/alias_suffix.py          |  7 +++++--
 tests/test_alias_suffixes.py | 22 ++++++++++++++++++++++
 2 files changed, 27 insertions(+), 2 deletions(-)

diff --git a/app/alias_suffix.py b/app/alias_suffix.py
index c8e0d4ea..40c54622 100644
--- a/app/alias_suffix.py
+++ b/app/alias_suffix.py
@@ -169,8 +169,11 @@ def get_alias_suffixes(
             alias_suffixes.insert(0, alias_suffix)
 
     if not default_domain_found:
-        sl_domain = SLDomain.get(user.default_alias_public_domain_id)
-        if sl_domain and not sl_domain.hidden:
+        domain_conditions = {"id": user.default_alias_public_domain_id, "hidden": False}
+        if not user.is_premium():
+            domain_conditions["premium_only"] = False
+        sl_domain = SLDomain.get_by(**domain_conditions)
+        if sl_domain:
             prefix = (
                 ""
                 if config.DISABLE_ALIAS_SUFFIX
diff --git a/tests/test_alias_suffixes.py b/tests/test_alias_suffixes.py
index 182f6c26..9706a13e 100644
--- a/tests/test_alias_suffixes.py
+++ b/tests/test_alias_suffixes.py
@@ -84,6 +84,28 @@ def test_get_default_domain_hidden():
     assert suffixes[0].domain != default_domain.domain
 
 
+def test_get_default_domain_is_premium_for_free_user():
+    user = create_new_user()
+    PartnerUser.create(
+        partner_id=get_proton_partner().id,
+        user_id=user.id,
+        external_user_id=random_token(10),
+        flush=True,
+    )
+    user.trial_end = None
+    default_domain = SLDomain.filter_by(partner_id=None, premium_only=True).first()
+    user.default_alias_public_domain_id = default_domain.id
+    Session.flush()
+    options = AliasOptions(
+        show_sl_domains=False, show_partner_domains=get_proton_partner()
+    )
+    suffixes = get_alias_suffixes(user, alias_options=options)
+    for suffix in suffixes:
+        domain = SLDomain.get_by(domain=suffix.domain)
+        assert not domain.premium_only
+    assert suffixes[0].domain != default_domain.domain
+
+
 def test_suffixes_are_valid():
     user = create_new_user()
     PartnerUser.create(