Fix proton integration issues (#1071)

* Fix proton integration issues

* Make external_user_id non nullable

* Fix tests

app/dashboard/views/setting.py

@@ -75,9 +75,6 @@ def get_proton_linked_account() -> Optional[str]:
     except ProtonPartnerNotSetUp:
         return None
 
-    if current_user.partner_id != proton_partner_id:
-        return None
-
     # It has. Retrieve the information for the PartnerUser
     proton_linked_account = PartnerUser.get_by(
         user_id=current_user.id, partner_id=proton_partner_id

app/models.py

@@ -3142,7 +3142,7 @@ class PartnerUser(Base, ModelMixin):
     partner_id = sa.Column(
         sa.ForeignKey("partner.id", ondelete="cascade"), nullable=False, index=True
     )
-    external_user_id = sa.Column(sa.String(128), unique=False, nullable=True)
+    external_user_id = sa.Column(sa.String(128), unique=False, nullable=False)
     partner_email = sa.Column(sa.String(255), unique=False, nullable=True)
 
     user = orm.relationship(User, foreign_keys=[user_id])

app/proton/proton_client.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from arrow import Arrow
 from dataclasses import dataclass
 from http import HTTPStatus
 from requests import Response, Session
@@ -96,7 +97,10 @@ def get_user(self) -> Optional[UserInformation]:
         if plan_value == PLAN_FREE:
             plan = SLPlan(type=SLPlanType.Free, expiration=None)
         elif plan_value == PLAN_PREMIUM:
-            plan = SLPlan(type=SLPlanType.Premium, expiration=info["PlanExpiration"])
+            plan = SLPlan(
+                type=SLPlanType.Premium,
+                expiration=Arrow.fromtimestamp(info["PlanExpiration"], tzinfo="utc"),
+            )
         else:
             raise Exception(f"Invalid value for plan: {plan_value}")
 

migrations/versions/2022_061016_36646e5dc6d9_make_external_user_id_non_nullable.py

@@ -0,0 +1,33 @@
+"""make external_user_id non nullable
+
+Revision ID: 36646e5dc6d9
+Revises: 82d3c7109ffb
+Create Date: 2022-06-10 16:07:11.538577
+
+"""
+import sqlalchemy_utils
+from alembic import op
+import sqlalchemy as sa
+
+
+# revision identifiers, used by Alembic.
+revision = '36646e5dc6d9'
+down_revision = '82d3c7109ffb'
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column('partner_user', 'external_user_id',
+               existing_type=sa.VARCHAR(length=128),
+               nullable=False)
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    op.alter_column('partner_user', 'external_user_id',
+               existing_type=sa.VARCHAR(length=128),
+               nullable=True)
+    # ### end Alembic commands ###

tests/models/test_partner_subscription.py

@@ -5,6 +5,7 @@
 
 
 def test_generate_partner_subscription(flask_client):
+    external_user_id = random_string()
     partner = Partner.create(
         name=random_string(10),
         contact_email=random_email(),
@@ -15,6 +16,7 @@ def test_generate_partner_subscription(flask_client):
         user_id=user.id,
         partner_id=partner.id,
         partner_email=random_email(),
+        external_user_id=external_user_id,
         commit=True,
     )