Increase retries to be ensured at most once a day

app/config.py

@@ -738,3 +738,5 @@ def read_hex_data(key: string, default: bytes) -> bytes:
 DROP_PGP_KEY_ATTACHMENTS_ON_REPLY = "DROP_PGP_KEY_ATTACHMENTS_ON_REPLY" in os.environ
 
 ENFORCE_OAUTH_CLIENT_APPROVED = "ENFORCE_OAUTH_CLIENT_APPROVED" in os.environ
+
+MAX_DOMAIN_CHECKS = max(int(os.environ.get("MAX_DOMAIN_CHECKS", 4)), 1)

tasks/check_custom_domains.py

@@ -67,16 +67,18 @@ def check_single_custom_domain(custom_domain: CustomDomain):
     if not is_mx_equivalent(mx_domains, expected_custom_domains):
         user = custom_domain.user
         LOG.w(
-            "The MX record is not correctly set for %s %s %s",
-            custom_domain,
-            user,
-            mx_domains,
+            f"The MX record is not correctly set for domain {custom_domain} of user {user}. Got {mx_domains}. Retried {custom_domain.nb_failed_checks} days",
         )
 
-        custom_domain.nb_failed_checks += 1
+        if (
+            not custom_domain.updated_at
+            or custom_domain.updated_at <= arrow.now().shift(days=-1)
+        ):
+            # Only update it once a day
+            custom_domain.nb_failed_checks += 1
 
-        # send alert if fail for 3 consecutive days
-        if custom_domain.nb_failed_checks > 2:
+        # send alert if fail for MAX_DOMAIN_CHECKS consecutive days
+        if custom_domain.nb_failed_checks > config.MAX_DOMAIN_CHECKS:
             domain_dns_url = f"{config.URL}/dashboard/domains/{custom_domain.id}/dns"
             LOG.w("Alert domain MX check fails %s about %s", user, custom_domain)
             send_email_with_rate_control(

tests/tasks/test_check_custom_domain.py

@@ -1,6 +1,7 @@
 import arrow
 from unittest.mock import patch
 
+from app.db import Session
 from app.models import CustomDomain
 from tasks.check_custom_domains import (
     check_all_custom_domains,
@@ -11,13 +12,18 @@
 
 def test_check_single_custom_domain_increments_failed_checks(flask_client):
     user = create_new_user()
+    # Set updated_at to 2 days ago so the increment logic runs
     custom_domain = CustomDomain.create(
         user_id=user.id,
         domain=random_string(),
         verified=True,
         nb_failed_checks=0,
         commit=True,
     )
+    # Manually set updated_at to simulate an old update
+    custom_domain.updated_at = arrow.now().shift(days=-2)
+    Session.commit()
+
     with patch("tasks.check_custom_domains.get_mx_domains", return_value=[]), patch(
         "tasks.check_custom_domains.is_mx_equivalent", return_value=False
     ), patch("tasks.check_custom_domains.send_email_with_rate_control"):
@@ -28,13 +34,18 @@ def test_check_single_custom_domain_increments_failed_checks(flask_client):
 
 def test_check_single_custom_domain_deactivates_after_threshold(flask_client):
     user = create_new_user()
+    # Start at 4 so that after one increment (4->5), 5 > MAX_DOMAIN_CHECKS (4) triggers deactivation
     custom_domain = CustomDomain.create(
         user_id=user.id,
         domain=random_string(),
         verified=True,
-        nb_failed_checks=2,
+        nb_failed_checks=4,
         commit=True,
     )
+    # Set updated_at to 2 days ago so the increment logic runs
+    custom_domain.updated_at = arrow.now().shift(days=-2)
+    Session.commit()
+
     with patch("tasks.check_custom_domains.get_mx_domains", return_value=[]), patch(
         "tasks.check_custom_domains.is_mx_equivalent", return_value=False
     ), patch("tasks.check_custom_domains.send_email_with_rate_control"):

tests/test_custom_domain_validation.py

@@ -296,6 +296,74 @@ def test_custom_domain_validation_validate_dkim_records_success():
     assert db_domain.dkim_verified is True
 
 
+def test_custom_domain_validation_validate_dkim_records_partner_domain_success():
+    """Test DKIM validation for partner domains with custom DKIM configuration."""
+    dkim_domain = random_domain()
+    partner_dkim_domain = random_domain()
+    dns_client = InMemoryDNSClient()
+
+    partner_id = get_proton_partner().id
+    validator = CustomDomainValidation(
+        dkim_domain, dns_client, partner_domains={partner_id: partner_dkim_domain}
+    )
+
+    user_domain = random_domain()
+    domain = create_custom_domain(user_domain)
+    domain.partner_id = partner_id
+    Session.commit()
+
+    # Set correct partner DKIM records (preferred)
+    dns_client.set_cname_record(
+        f"dkim._domainkey.{user_domain}", f"dkim._domainkey.{partner_dkim_domain}"
+    )
+    dns_client.set_cname_record(
+        f"dkim02._domainkey.{user_domain}", f"dkim02._domainkey.{partner_dkim_domain}"
+    )
+    dns_client.set_cname_record(
+        f"dkim03._domainkey.{user_domain}", f"dkim03._domainkey.{partner_dkim_domain}"
+    )
+
+    res = validator.validate_dkim_records(domain)
+    assert len(res) == 0
+
+    db_domain = CustomDomain.get_by(id=domain.id)
+    assert db_domain.dkim_verified is True
+
+
+def test_custom_domain_validation_validate_dkim_records_partner_domain_fallback_success():
+    """Test DKIM validation for partner domains falls back to SL domain if partner DKIM not set."""
+    dkim_domain = random_domain()
+    partner_mx_domain = random_domain()  # Partner has MX config but not DKIM
+    dns_client = InMemoryDNSClient()
+
+    partner_id = get_proton_partner().id
+    validator = CustomDomainValidation(
+        dkim_domain, dns_client, partner_domains={partner_id: partner_mx_domain}
+    )
+
+    user_domain = random_domain()
+    domain = create_custom_domain(user_domain)
+    domain.partner_id = partner_id
+    Session.commit()
+
+    # Set correct SL DKIM records (fallback)
+    dns_client.set_cname_record(
+        f"dkim._domainkey.{user_domain}", f"dkim._domainkey.{dkim_domain}"
+    )
+    dns_client.set_cname_record(
+        f"dkim02._domainkey.{user_domain}", f"dkim02._domainkey.{dkim_domain}"
+    )
+    dns_client.set_cname_record(
+        f"dkim03._domainkey.{user_domain}", f"dkim03._domainkey.{dkim_domain}"
+    )
+
+    res = validator.validate_dkim_records(domain)
+    assert len(res) == 0
+
+    db_domain = CustomDomain.get_by(id=domain.id)
+    assert db_domain.dkim_verified is True
+
+
 # validate_ownership
 def test_custom_domain_validation_validate_ownership_empty_records_failure():
     dns_client = InMemoryDNSClient()
@@ -428,6 +496,39 @@ def test_custom_domain_validation_validate_mx_records_success():
     assert db_domain.verified is True
 
 
+def test_custom_domain_validation_validate_mx_records_partner_domain_success():
+    """Test MX validation for partner domains with custom MX configuration."""
+    dns_client = InMemoryDNSClient()
+    partner_id = get_proton_partner().id
+    partner_mx_domain = random_domain()
+
+    validator = CustomDomainValidation(
+        random_domain(), dns_client, partner_domains={partner_id: partner_mx_domain}
+    )
+
+    domain = create_custom_domain(random_domain())
+    domain.partner_id = partner_id
+    Session.commit()
+
+    # Get expected records (should include partner MX + fallback SL MX)
+    expected_records = validator.get_expected_mx_records(domain)
+
+    # Set MX records with partner MX first (preferred priority)
+    dns_records = {}
+    for priority in expected_records:
+        # Use the recommended record (partner MX)
+        dns_records[priority] = [expected_records[priority].recommended]
+
+    dns_client.set_mx_records(domain.domain, dns_records)
+    res = validator.validate_mx_records(domain)
+
+    assert res.success is True
+    assert len(res.errors) == 0
+
+    db_domain = CustomDomain.get_by(id=domain.id)
+    assert db_domain.verified is True
+
+
 # validate_spf_records
 def test_custom_domain_validation_validate_spf_records_empty_failure():
     dns_client = InMemoryDNSClient()