Remove get_proton_partner_id function

app/dashboard/views/setting.py

@@ -48,7 +48,7 @@ from app.models import (
     AppleSubscription,
     PartnerUser,
 )
-from app.proton.proton_callback_handler import get_proton_partner_id
+from app.proton.proton_callback_handler import get_proton_partner
 from app.utils import random_string, sanitize_email
 
 
@@ -70,7 +70,7 @@ class PromoCodeForm(FlaskForm):
 def get_proton_linked_account() -> Optional[str]:
     # Check if the current user has a partner_id
     try:
-        proton_partner_id = get_proton_partner_id()
+        proton_partner_id = get_proton_partner().id
     except ProtonPartnerNotSetUp:
         return None
 
@@ -444,7 +444,7 @@ def unlink_proton_account():
     current_user.partner_id = None
     current_user.partner_user_id = None
     partner_user = PartnerUser.get_by(
-        user_id=current_user.id, partner_id=get_proton_partner_id()
+        user_id=current_user.id, partner_id=get_proton_partner().id
     )
     if partner_user is not None:
         PartnerUser.delete(partner_user.id)

app/proton/proton_callback_handler.py

@@ -25,10 +25,6 @@ def get_proton_partner() -> Partner:
     return _PROTON_PARTNER
 
 
-def get_proton_partner_id() -> int:
-    return get_proton_partner().id
-
-
 class Action(enum.Enum):
     Login = 1
     Link = 2

tests/proton/test_proton_callback_handler.py

@@ -5,7 +5,6 @@ from app.proton.proton_client import ProtonClient, UserInformation, ProtonPlan
 from app.proton.proton_callback_handler import (
     ProtonCallbackHandler,
     get_proton_partner,
-    get_proton_partner_id,
     get_login_strategy,
     process_link_case,
     ProtonUser,
@@ -66,11 +65,11 @@ def create_user(email: str = None) -> User:
 def create_user_for_partner(partner_user_id: str, email: str = None) -> User:
     email = email if email is not None else random_email()
     user = User.create(email=email)
-    user.partner_id = get_proton_partner_id()
+    user.partner_id = get_proton_partner().id
     user.partner_user_id = partner_user_id
 
     PartnerUser.create(
-        user_id=user.id, partner_id=get_proton_partner_id(), partner_email=email
+        user_id=user.id, partner_id=get_proton_partner().id, partner_email=email
     )
     Session.commit()
     return user
@@ -108,7 +107,7 @@ def test_proton_callback_handler_existant_sl_user():
     assert res.user is not None
     assert res.user.id == sl_user.id
 
-    sa = PartnerUser.get_by(user_id=sl_user.id, partner_id=get_proton_partner_id())
+    sa = PartnerUser.get_by(user_id=sl_user.id, partner_id=get_proton_partner().id)
     assert sa is not None
     assert sa.partner_email == user.email
 
@@ -191,7 +190,7 @@ def test_link_account_with_proton_account_same_address(flask_client):
     assert res.flash_message is not None
 
     updated_user = User.get(sl_user.id)
-    assert updated_user.partner_id == get_proton_partner_id()
+    assert updated_user.partner_id == get_proton_partner().id
     assert updated_user.partner_user_id == proton_user_id
 
 
@@ -211,7 +210,7 @@ def test_link_account_with_proton_account_different_address(flask_client):
     assert res.flash_message is not None
 
     updated_user = User.get(sl_user.id)
-    assert updated_user.partner_id == get_proton_partner_id()
+    assert updated_user.partner_id == get_proton_partner().id
     assert updated_user.partner_user_id == proton_user_id
 
 
@@ -238,7 +237,7 @@ def test_link_account_with_proton_account_same_address_but_linked_to_other_user(
     assert res.flash_message is not None
 
     updated_user_1 = User.get(sl_user_1.id)
-    assert updated_user_1.partner_id == get_proton_partner_id()
+    assert updated_user_1.partner_id == get_proton_partner().id
     assert updated_user_1.partner_user_id == proton_user_id
 
     updated_user_2 = User.get(sl_user_2.id)
@@ -268,10 +267,10 @@ def test_link_account_with_proton_account_different_address_and_linked_to_other_
     assert res.flash_message is not None
 
     updated_user_1 = User.get(sl_user_1.id)
-    assert updated_user_1.partner_id == get_proton_partner_id()
+    assert updated_user_1.partner_id == get_proton_partner().id
     assert updated_user_1.partner_user_id == proton_user_id
     partner_user_1 = PartnerUser.get_by(
-        user_id=sl_user_1.id, partner_id=get_proton_partner_id()
+        user_id=sl_user_1.id, partner_id=get_proton_partner().id
     )
     assert partner_user_1 is not None
     assert partner_user_1.partner_email == proton_user.email
@@ -280,7 +279,7 @@ def test_link_account_with_proton_account_different_address_and_linked_to_other_
     assert updated_user_2.partner_id is None
     assert updated_user_2.partner_user_id is None
     partner_user_2 = PartnerUser.get_by(
-        user_id=sl_user_2.id, partner_id=get_proton_partner_id()
+        user_id=sl_user_2.id, partner_id=get_proton_partner().id
     )
     assert partner_user_2 is None