Fix: Always include default domain in the list of domains (#1768)

* Fix: Always include default domain in the list of domains

* Add premium test

---------

Co-authored-by: Adrià Casajús <adria.casajus@proton.ch>

app/models.py

@@ -1016,6 +1016,10 @@ class User(Base, ModelMixin, UserMixin, PasswordOracle):
         if not self.is_premium():
             conditions.append(SLDomain.premium_only == False)  # noqa: E712
         partner_domain_cond = []  # noqa:E711
+        if self.default_alias_public_domain_id is not None:
+            partner_domain_cond.append(
+                SLDomain.id == self.default_alias_public_domain_id
+            )
         if alias_options.show_partner_domains is not None:
             partner_user = PartnerUser.filter_by(
                 user_id=self.id, partner_id=alias_options.show_partner_domains.id

tests/dashboard/test_custom_alias.py

@@ -316,6 +316,10 @@ def test_add_alias_in_global_trash(flask_client):
 def test_add_alias_in_custom_domain_trash(flask_client):
     user = login(flask_client)
 
+    for deleted_domain in DomainDeletedAlias.all():
+        Session.delete(deleted_domain)
+    Session.flush()
+
     domain = random_domain()
     custom_domain = CustomDomain.create(
         user_id=user.id, domain=domain, ownership_verified=True, commit=True

tests/test_domains.py

@@ -128,3 +128,74 @@ def test_get_premium_with_partner_domains():
     assert [d.domain for d in domains] == user.available_sl_domains(
         alias_options=options
     )
+
+
+def test_get_partner_and_free_default_domain():
+    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.default_alias_public_domain_id = (
+        SLDomain.filter_by(partner_id=None, hidden=False).first().id
+    )
+    Session.flush()
+    options = AliasOptions(
+        show_sl_domains=False, show_partner_domains=get_proton_partner()
+    )
+    domains = user.get_sl_domains(alias_options=options)
+    assert len(domains) == 3
+    assert domains[0].domain == "premium_partner"
+    assert domains[1].domain == "free_partner"
+    assert domains[2].domain == "free_non_partner"
+    assert [d.domain for d in domains] == user.available_sl_domains(
+        alias_options=options
+    )
+
+
+def test_get_free_partner_and_premium_default_domain():
+    user = create_new_user()
+    user.trial_end = None
+    PartnerUser.create(
+        partner_id=get_proton_partner().id,
+        user_id=user.id,
+        external_user_id=random_token(10),
+        flush=True,
+    )
+    user.default_alias_public_domain_id = (
+        SLDomain.filter_by(partner_id=None, hidden=False, premium_only=True).first().id
+    )
+    Session.flush()
+    options = AliasOptions(
+        show_sl_domains=False, show_partner_domains=get_proton_partner()
+    )
+    domains = user.get_sl_domains(alias_options=options)
+    assert len(domains) == 1
+    assert domains[0].domain == "free_partner"
+    assert [d.domain for d in domains] == user.available_sl_domains(
+        alias_options=options
+    )
+
+
+def test_get_free_partner_and_hidden_default_domain():
+    user = create_new_user()
+    user.trial_end = None
+    PartnerUser.create(
+        partner_id=get_proton_partner().id,
+        user_id=user.id,
+        external_user_id=random_token(10),
+        flush=True,
+    )
+    user.default_alias_public_domain_id = SLDomain.filter_by(hidden=True).first().id
+    Session.flush()
+    options = AliasOptions(
+        show_sl_domains=False, show_partner_domains=get_proton_partner()
+    )
+    domains = user.get_sl_domains(alias_options=options)
+    assert len(domains) == 1
+    assert domains[0].domain == "free_partner"
+    assert [d.domain for d in domains] == user.available_sl_domains(
+        alias_options=options
+    )