From b720dfc381fc3aa728ffd1bae6181ef2afe69dd9 Mon Sep 17 00:00:00 2001
From: Son <nguyenkims@users.noreply.github.com>
Date: Thu, 5 May 2022 15:05:39 +0200
Subject: [PATCH] allow the code to run without proton partner

---
 app/dashboard/views/setting.py        | 7 ++++++-
 app/errors.py                         | 4 ++++
 app/proton/proton_callback_handler.py | 3 ++-
 3 files changed, 12 insertions(+), 2 deletions(-)

diff --git a/app/dashboard/views/setting.py b/app/dashboard/views/setting.py
index f5b382d1..4a768874 100644
--- a/app/dashboard/views/setting.py
+++ b/app/dashboard/views/setting.py
@@ -28,6 +28,7 @@ from app.email_utils import (
     email_can_be_used_as_mailbox,
     personal_email_already_used,
 )
+from app.errors import ProtonPartnerNotSetUp
 from app.log import LOG
 from app.models import (
     BlockBehaviourEnum,
@@ -68,7 +69,11 @@ class PromoCodeForm(FlaskForm):
 
 def get_proton_linked_account() -> Optional[str]:
     # Check if the current user has a partner_id
-    proton_partner_id = get_proton_partner_id()
+    try:
+        proton_partner_id = get_proton_partner_id()
+    except ProtonPartnerNotSetUp:
+        return None
+
     if current_user.partner_id != proton_partner_id:
         return None
 
diff --git a/app/errors.py b/app/errors.py
index 26be6091..f75e83d1 100644
--- a/app/errors.py
+++ b/app/errors.py
@@ -56,3 +56,7 @@ class MailSentFromReverseAlias(SLException):
     """raised when receiving an email sent from a reverse alias"""
 
     pass
+
+
+class ProtonPartnerNotSetUp(SLException):
+    pass
diff --git a/app/proton/proton_callback_handler.py b/app/proton/proton_callback_handler.py
index abdecadf..e444ed16 100644
--- a/app/proton/proton_callback_handler.py
+++ b/app/proton/proton_callback_handler.py
@@ -5,6 +5,7 @@ from flask import url_for
 from typing import Optional
 
 from app.db import Session
+from app.errors import ProtonPartnerNotSetUp
 from app.models import User, PartnerUser, Partner
 from app.proton.proton_client import ProtonClient, ProtonUser
 from app.utils import random_string
@@ -18,7 +19,7 @@ def get_proton_partner_id() -> int:
     if _PROTON_PARTNER_ID is None:
         partner = Partner.get_by(name=PROTON_PARTNER_NAME)
         if partner is None:
-            raise Exception("Could not find Proton Partner instance")
+            raise ProtonPartnerNotSetUp
         _PROTON_PARTNER_ID = partner.id
 
     return _PROTON_PARTNER_ID