4.24.0

app/tests/models/test_user.py

@@ -1,13 +1,7 @@
 from app import config
 from app.db import Session
 from app.models import User, Job
-from tests.utils import create_new_user, random_email
-
-
-def test_available_sl_domains(flask_client):
-    user = create_new_user()
-
-    assert set(user.available_sl_domains()) == {"d1.test", "d2.test", "sl.local"}
+from tests.utils import random_email
 
 
 def test_create_from_partner(flask_client):

app/tests/test_domains.py

@@ -0,0 +1,130 @@
+from app.db import Session
+from app.models import SLDomain, PartnerUser, AliasOptions
+from app.proton.utils import get_proton_partner
+from init_app import add_sl_domains
+from tests.utils import create_new_user, random_token
+
+
+def setup_module():
+    Session.query(SLDomain).delete()
+    SLDomain.create(
+        domain="hidden", premium_only=False, flush=True, order=5, hidden=True
+    )
+    SLDomain.create(domain="free_non_partner", premium_only=False, flush=True, order=4)
+    SLDomain.create(
+        domain="premium_non_partner", premium_only=True, flush=True, order=3
+    )
+    SLDomain.create(
+        domain="free_partner",
+        premium_only=False,
+        flush=True,
+        partner_id=get_proton_partner().id,
+        order=2,
+    )
+    SLDomain.create(
+        domain="premium_partner",
+        premium_only=True,
+        flush=True,
+        partner_id=get_proton_partner().id,
+        order=1,
+    )
+    Session.commit()
+
+
+def teardown_module():
+    Session.query(SLDomain).delete()
+    add_sl_domains()
+
+
+def test_get_non_partner_domains():
+    user = create_new_user()
+    domains = user.get_sl_domains()
+    # Premium
+    assert len(domains) == 2
+    assert domains[0].domain == "premium_non_partner"
+    assert domains[1].domain == "free_non_partner"
+    assert [d.domain for d in domains] == user.available_sl_domains()
+    # Free
+    user.trial_end = None
+    Session.flush()
+    domains = user.get_sl_domains()
+    assert len(domains) == 1
+    assert domains[0].domain == "free_non_partner"
+    assert [d.domain for d in domains] == user.available_sl_domains()
+
+
+def test_get_free_with_partner_domains():
+    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,
+    )
+    domains = user.get_sl_domains()
+    # Default
+    assert len(domains) == 1
+    assert domains[0].domain == "free_non_partner"
+    assert [d.domain for d in domains] == user.available_sl_domains()
+    # Show partner domains
+    options = AliasOptions(
+        show_sl_domains=True, show_partner_domains=get_proton_partner()
+    )
+    domains = user.get_sl_domains(alias_options=options)
+    assert len(domains) == 2
+    assert domains[0].domain == "free_partner"
+    assert domains[1].domain == "free_non_partner"
+    assert [d.domain for d in domains] == user.available_sl_domains(
+        alias_options=options
+    )
+    # Only partner domains
+    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_premium_with_partner_domains():
+    user = create_new_user()
+    PartnerUser.create(
+        partner_id=get_proton_partner().id,
+        user_id=user.id,
+        external_user_id=random_token(10),
+        flush=True,
+    )
+    domains = user.get_sl_domains()
+    # Default
+    assert len(domains) == 2
+    assert domains[0].domain == "premium_non_partner"
+    assert domains[1].domain == "free_non_partner"
+    assert [d.domain for d in domains] == user.available_sl_domains()
+    # Show partner domains
+    options = AliasOptions(
+        show_sl_domains=True, show_partner_domains=get_proton_partner()
+    )
+    domains = user.get_sl_domains(alias_options=options)
+    assert len(domains) == 4
+    assert domains[0].domain == "premium_partner"
+    assert domains[1].domain == "free_partner"
+    assert domains[2].domain == "premium_non_partner"
+    assert domains[3].domain == "free_non_partner"
+    assert [d.domain for d in domains] == user.available_sl_domains(
+        alias_options=options
+    )
+    # Only partner domains
+    options = AliasOptions(
+        show_sl_domains=False, show_partner_domains=get_proton_partner()
+    )
+    domains = user.get_sl_domains(alias_options=options)
+    assert len(domains) == 2
+    assert domains[0].domain == "premium_partner"
+    assert domains[1].domain == "free_partner"
+    assert [d.domain for d in domains] == user.available_sl_domains(
+        alias_options=options
+    )

app/tests/test_utils.py

@@ -13,7 +13,7 @@ def test_random_words():
     assert s.count("_") == 1
     assert len(s) > 3
     s = random_words(2, 3)
-    assert s.count("_") == 0
+    assert s.count("_") == 1
     assert s[-1] in (str(i) for i in range(10))