Allow to create users from partner (#1095)

* Allow to create users from partner

* Fix tests

* Update tests/test_account_linking.py

Co-authored-by: Adrià Casajús <acasajus@users.noreply.github.com>

* Fix lint

Co-authored-by: Adrià Casajús <acasajus@users.noreply.github.com>

app/account_linking.py

@@ -29,6 +29,7 @@ class PartnerLinkRequest:
     email: str
     external_user_id: str
     plan: SLPlan
+    from_partner: bool
 
 
 @dataclass
@@ -121,6 +122,7 @@ class NewUserStrategy(ClientMergeStrategy):
             email=self.link_request.email,
             name=self.link_request.name,
             password=random_string(20),
+            from_partner=self.link_request.from_partner,
         )
         partner_user = PartnerUser.create(
             user_id=new_user.id,
@@ -145,7 +147,7 @@ class NewUserStrategy(ClientMergeStrategy):
         )
 
 
-class ExistingUnlinedUserStrategy(ClientMergeStrategy):
+class ExistingUnlinkedUserStrategy(ClientMergeStrategy):
     def process(self) -> LinkResult:
 
         partner_user = ensure_partner_user_exists_for_user(
@@ -175,7 +177,7 @@ def get_login_strategy(
     if other_partner_user is not None:
         return LinkedWithAnotherPartnerUserStrategy(link_request, user, partner)
     # There is a SimpleLogin user with the partner_user's e-mail
-    return ExistingUnlinedUserStrategy(link_request, user, partner)
+    return ExistingUnlinkedUserStrategy(link_request, user, partner)
 
 
 def process_login_case(

app/models.py

@@ -302,7 +302,8 @@ class Fido(Base, ModelMixin):
 class User(Base, ModelMixin, UserMixin, PasswordOracle):
     __tablename__ = "users"
 
-    FLAG_FREE_DISABLE_CREATE_ALIAS = 1
+    FLAG_FREE_DISABLE_CREATE_ALIAS = 1 << 0
+    FLAG_CREATED_FROM_PARTNER = 1 << 1
 
     email = sa.Column(sa.String(256), unique=True, nullable=False)
 
@@ -528,7 +529,7 @@ class User(Base, ModelMixin, UserMixin, PasswordOracle):
             return str(self.id)
 
     @classmethod
-    def create(cls, email, name="", password=None, **kwargs):
+    def create(cls, email, name="", password=None, from_partner=False, **kwargs):
         user: User = super(User, cls).create(email=email, name=name, **kwargs)
 
         if password:
@@ -557,6 +558,15 @@ class User(Base, ModelMixin, UserMixin, PasswordOracle):
         if "alternative_id" not in kwargs:
             user.alternative_id = str(uuid.uuid4())
 
+        # If the user is created from partner, do not notify
+        # nor give a trial
+        if from_partner:
+            user.flags = User.FLAG_CREATED_FROM_PARTNER
+            user.notification = False
+            user.trial_end = None
+            Session.flush()
+            return user
+
         if DISABLE_ONBOARDING:
             LOG.d("Disable onboarding emails")
             return user

app/proton/proton_callback_handler.py

@@ -113,6 +113,7 @@ class ProtonCallbackHandler:
             external_user_id=proton_user.id,
             name=proton_user.name,
             plan=proton_user.plan,
+            from_partner=False,  # The user has started this flow, so we don't mark it as created by a partner
         )
 
     def __get_proton_user(self) -> Optional[ProtonUser]:

tests/models/test_user.py

@@ -1,7 +1,17 @@
-from tests.utils import create_new_user
+from app.models import User
+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"}
+
+
+def test_create_from_partner(flask_client):
+    user = User.create(email=random_email(), from_partner=True)
+    assert User.FLAG_CREATED_FROM_PARTNER == (
+        user.flags & User.FLAG_CREATED_FROM_PARTNER
+    )
+    assert user.notification is False
+    assert user.trial_end is None

tests/proton/test_proton_callback_handler.py

@@ -44,6 +44,12 @@ def test_proton_callback_handler_unexistant_sl_user():
     assert res.user is not None
     assert res.user.email == email
     assert res.user.name == name
+    # Ensure the user is not marked as created from partner
+    assert User.FLAG_CREATED_FROM_PARTNER != (
+        res.user.flags & User.FLAG_CREATED_FROM_PARTNER
+    )
+    assert res.user.notification is True
+    assert res.user.trial_end is not None
 
     partner_user = PartnerUser.get_by(
         partner_id=get_proton_partner().id, user_id=res.user.id
@@ -68,6 +74,12 @@ def test_proton_callback_handler_existant_sl_user():
 
     assert res.user is not None
     assert res.user.id == sl_user.id
+    # Ensure the user is not marked as created from partner
+    assert User.FLAG_CREATED_FROM_PARTNER != (
+        res.user.flags & User.FLAG_CREATED_FROM_PARTNER
+    )
+    assert res.user.notification is True
+    assert res.user.trial_end is not None
 
     sa = PartnerUser.get_by(user_id=sl_user.id, partner_id=get_proton_partner().id)
     assert sa is not None

tests/test_account_linking.py

@@ -3,10 +3,11 @@ from arrow import Arrow
 
 from app.account_linking import (
     process_link_case,
+    process_login_case,
     get_login_strategy,
     ensure_partner_user_exists_for_user,
     NewUserStrategy,
-    ExistingUnlinedUserStrategy,
+    ExistingUnlinkedUserStrategy,
     LinkedWithAnotherPartnerUserStrategy,
     SLPlan,
     SLPlanType,
@@ -27,6 +28,7 @@ def random_link_request(
     name: str = None,
     email: str = None,
     plan: SLPlan = None,
+    from_partner: bool = False,
 ) -> PartnerLinkRequest:
     external_user_id = (
         external_user_id if external_user_id is not None else random_string()
@@ -39,6 +41,7 @@ def random_link_request(
         email=email,
         external_user_id=external_user_id,
         plan=SLPlan(type=plan, expiration=Arrow.utcnow().shift(hours=2)),
+        from_partner=from_partner,
     )
 
 
@@ -72,6 +75,38 @@ def test_get_strategy_unexistant_sl_user():
     assert isinstance(strategy, NewUserStrategy)
 
 
+def test_login_case_from_partner():
+    partner = get_proton_partner()
+    res = process_login_case(
+        random_link_request(
+            external_user_id=random_string(),
+            from_partner=True,
+        ),
+        partner,
+    )
+
+    assert res.strategy == NewUserStrategy.__name__
+    assert res.user is not None
+    assert User.FLAG_CREATED_FROM_PARTNER == (
+        res.user.flags & User.FLAG_CREATED_FROM_PARTNER
+    )
+
+
+def test_login_case_from_web():
+    partner = get_proton_partner()
+    res = process_login_case(
+        random_link_request(
+            external_user_id=random_string(),
+            from_partner=False,
+        ),
+        partner,
+    )
+
+    assert res.strategy == NewUserStrategy.__name__
+    assert res.user is not None
+    assert 0 == (res.user.flags & User.FLAG_CREATED_FROM_PARTNER)
+
+
 def test_get_strategy_existing_sl_user():
     email = random_email()
     user = User.create(email, commit=True)
@@ -80,7 +115,7 @@ def test_get_strategy_existing_sl_user():
         user=user,
         partner=get_proton_partner(),
     )
-    assert isinstance(strategy, ExistingUnlinedUserStrategy)
+    assert isinstance(strategy, ExistingUnlinkedUserStrategy)
 
 
 def test_get_strategy_existing_sl_user_linked_with_different_proton_account():
@@ -266,6 +301,7 @@ def test_ensure_partner_user_exists_for_user_raises_exception_when_linked_to_ano
             email=user_email,
             external_user_id=external_id_1,
             plan=SLPlan(type=SLPlanType.Free, expiration=None),
+            from_partner=False,
         ),
         user,
         partner_1,
@@ -279,6 +315,7 @@ def test_ensure_partner_user_exists_for_user_raises_exception_when_linked_to_ano
                 email=user_email,
                 external_user_id=external_id_2,
                 plan=SLPlan(type=SLPlanType.Free, expiration=None),
+                from_partner=False,
             ),
             user,
             partner_2,