From f940df0b1d9e5563e5b1c225e52ac2e55a48bd50 Mon Sep 17 00:00:00 2001
From: Carlos Quintana <carlos.quintana@proton.ch>
Date: Mon, 16 Sep 2024 16:02:54 +0200
Subject: [PATCH] chore: allow to specify partner_id to custom_domain

---
 app/custom_domain_utils.py        |  8 +++++++-
 tests/test_custom_domain_utils.py | 32 +++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 1 deletion(-)

diff --git a/app/custom_domain_utils.py b/app/custom_domain_utils.py
index 5a1608eca..30e891fba 100644
--- a/app/custom_domain_utils.py
+++ b/app/custom_domain_utils.py
@@ -65,7 +65,9 @@ def can_domain_be_used(user: User, domain: str) -> Optional[str]:
         return None
 
 
-def create_custom_domain(user: User, domain: str) -> CreateCustomDomainResult:
+def create_custom_domain(
+    user: User, domain: str, partner_id: Optional[int] = None
+) -> CreateCustomDomainResult:
     if not user.is_premium():
         return CreateCustomDomainResult(
             message="Only premium plan can add custom domain",
@@ -91,6 +93,10 @@ def create_custom_domain(user: User, domain: str) -> CreateCustomDomainResult:
             )
             new_custom_domain.ownership_verified = True
 
+    # Add the partner_id in case it's passed
+    if partner_id is not None:
+        new_custom_domain.partner_id = partner_id
+
     Session.commit()
 
     return CreateCustomDomainResult(
diff --git a/tests/test_custom_domain_utils.py b/tests/test_custom_domain_utils.py
index 0cf1db46a..07adc7e04 100644
--- a/tests/test_custom_domain_utils.py
+++ b/tests/test_custom_domain_utils.py
@@ -10,6 +10,7 @@
 )
 from app.db import Session
 from app.models import User, CustomDomain, Mailbox
+from tests.utils import get_proton_partner
 from tests.utils import create_new_user, random_string, random_domain
 
 user: Optional[User] = None
@@ -108,3 +109,34 @@ def test_can_create_custom_domain():
 
     assert res.instance.domain == domain
     assert res.instance.user_id == user.id
+
+
+def test_can_create_custom_domain_validates_if_parent_is_validated():
+    root_domain = random_domain()
+    subdomain = f"{random_string(10)}.{root_domain}"
+
+    # Create custom domain with the root domain
+    CustomDomain.create(
+        user_id=user.id,
+        domain=root_domain,
+        verified=True,
+        ownership_verified=True,
+        commit=True,
+    )
+
+    # Create custom domain with subdomain. Should automatically be verified
+    res = create_custom_domain(user=user, domain=subdomain)
+    assert res.success is True
+    assert res.instance.domain == subdomain
+    assert res.instance.user_id == user.id
+    assert res.instance.ownership_verified is True
+
+
+def test_creates_custom_domain_with_partner_id():
+    domain = random_domain()
+    proton_partner = get_proton_partner()
+    res = create_custom_domain(user=user, domain=domain, partner_id=proton_partner.id)
+    assert res.success is True
+    assert res.instance.domain == domain
+    assert res.instance.user_id == user.id
+    assert res.instance.partner_id == proton_partner.id