4.65.3

app/tests/models/test_mailbox.py

@@ -0,0 +1,37 @@
+from app import config
+from app.dns_utils import set_global_dns_client, InMemoryDNSClient
+from app.email_utils import get_email_local_part
+from app.models import Mailbox
+from tests.utils import create_new_user, random_email
+
+dns_client = InMemoryDNSClient()
+
+
+def setup_module():
+    set_global_dns_client(dns_client)
+
+
+def teardown_module():
+    set_global_dns_client(None)
+
+
+def test_is_proton_with_email_domain():
+    user = create_new_user()
+    mailbox = Mailbox.create(
+        user_id=user.id, email=f"test@{config.PROTON_EMAIL_DOMAINS[0]}"
+    )
+    assert mailbox.is_proton()
+    mailbox = Mailbox.create(user_id=user.id, email="a@b.c")
+    assert not mailbox.is_proton()
+
+
+def test_is_proton_with_mx_domain():
+    email = random_email()
+    dns_client.set_mx_records(
+        get_email_local_part(email), {10: config.PROTON_MX_SERVERS}
+    )
+    user = create_new_user()
+    mailbox = Mailbox.create(user_id=user.id, email=email)
+    assert mailbox.is_proton()
+    dns_client.set_mx_records(get_email_local_part(email), {10: ["nowhere.net"]})
+    assert not mailbox.is_proton()

app/tests/models/test_user.py

@@ -2,7 +2,7 @@ import arrow
 from app import config
 from app.db import Session
 from app.models import User, Job, PartnerSubscription, PartnerUser, ManualSubscription
-from app.proton.utils import get_proton_partner
+from app.proton.proton_partner import get_proton_partner
 from tests.utils import random_email, random_token