From 56ec95bc93bc0e8fee48f2724733d62f41fac324 Mon Sep 17 00:00:00 2001
From: Carlos Quintana <74399022+cquintana92@users.noreply.github.com>
Date: Fri, 10 Jun 2022 16:21:56 +0200
Subject: [PATCH] Fix proton integration issues (#1071)

* Fix proton integration issues

* Make external_user_id non nullable

* Fix tests
---
 app/dashboard/views/setting.py                |  3 --
 app/models.py                                 |  2 +-
 app/proton/proton_client.py                   |  6 +++-
 ...c6d9_make_external_user_id_non_nullable.py | 33 +++++++++++++++++++
 tests/models/test_partner_subscription.py     |  2 ++
 5 files changed, 41 insertions(+), 5 deletions(-)
 create mode 100644 migrations/versions/2022_061016_36646e5dc6d9_make_external_user_id_non_nullable.py

diff --git a/app/dashboard/views/setting.py b/app/dashboard/views/setting.py
index 2b3ee9ca..6811e221 100644
--- a/app/dashboard/views/setting.py
+++ b/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
diff --git a/app/models.py b/app/models.py
index 9797c37d..2d50dd21 100644
--- a/app/models.py
+++ b/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])
diff --git a/app/proton/proton_client.py b/app/proton/proton_client.py
index 331862a5..82282d0d 100644
--- a/app/proton/proton_client.py
+++ b/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 @@ class HttpProtonClient(ProtonClient):
         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}")
 
diff --git a/migrations/versions/2022_061016_36646e5dc6d9_make_external_user_id_non_nullable.py b/migrations/versions/2022_061016_36646e5dc6d9_make_external_user_id_non_nullable.py
new file mode 100644
index 00000000..52295a13
--- /dev/null
+++ b/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 ###
diff --git a/tests/models/test_partner_subscription.py b/tests/models/test_partner_subscription.py
index aa1c376c..bf400076 100644
--- a/tests/models/test_partner_subscription.py
+++ b/tests/models/test_partner_subscription.py
@@ -5,6 +5,7 @@ from tests.utils import create_new_user, random_email
 
 
 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,
     )