fix(socialaccount): STORE_TOKENS while SocialApp not in db

pennersr · Jan 18, 2024 · 5f71a37 · 5f71a37
1 parent 3064115
commit 5f71a37
Show file tree

Hide file tree

Showing 2 changed files with 46 additions and 20 deletions.
diff --git a/allauth/socialaccount/models.py b/allauth/socialaccount/models.py
@@ -304,28 +304,35 @@ def _lookup_by_socialaccount(self):
             signals.social_account_updated.send(
                 sender=SocialLogin, request=context.request, sociallogin=self
             )
-            # Update token
-            if app_settings.STORE_TOKENS and self.token:
-                assert not self.token.pk
-                try:
-                    t = SocialToken.objects.get(
-                        account=self.account, app=self.token.app
-                    )
-                    t.token = self.token.token
-                    if self.token.token_secret:
-                        # only update the refresh token if we got one
-                        # many oauth2 providers do not resend the refresh token
-                        t.token_secret = self.token.token_secret
-                    t.expires_at = self.token.expires_at
-                    t.save()
-                    self.token = t
-                except SocialToken.DoesNotExist:
-                    self.token.account = a
-                    self.token.save()
+            self._store_token()
             return True
         except SocialAccount.DoesNotExist:
             pass
 
+    def _store_token(self):
+        # Update token
+        if not app_settings.STORE_TOKENS or not self.token:
+            return
+        assert not self.token.pk
+        app = self.token.app
+        if app and not app.pk:
+            # If the app is not stored in the db, leave the FK empty.
+            app = None
+        try:
+            t = SocialToken.objects.get(account=self.account, app=app)
+            t.token = self.token.token
+            if self.token.token_secret:
+                # only update the refresh token if we got one
+                # many oauth2 providers do not resend the refresh token
+                t.token_secret = self.token.token_secret
+            t.expires_at = self.token.expires_at
+            t.save()
+            self.token = t
+        except SocialToken.DoesNotExist:
+            self.token.account = self.account
+            self.token.app = app
+            self.token.save()
+
     def _lookup_by_email(self):
         emails = [e.email for e in self.email_addresses if e.verified]
         for email in emails:

diff --git a/allauth/socialaccount/tests/test_login.py b/allauth/socialaccount/tests/test_login.py
@@ -13,7 +13,7 @@
 from allauth.account.authentication import AUTHENTICATION_METHODS_SESSION_KEY
 from allauth.core import context
 from allauth.socialaccount.helpers import complete_social_login
-from allauth.socialaccount.models import SocialAccount
+from allauth.socialaccount.models import SocialAccount, SocialToken
 from allauth.socialaccount.providers.base import AuthProcess
 
 
@@ -92,6 +92,7 @@ def test_login_cancelled(client):
     assertTemplateUsed(resp, "socialaccount/login_cancelled.html")
 
 
+@pytest.mark.parametrize("store_tokens", [False, True])
 @pytest.mark.parametrize(
     "process,did_record",
     [
@@ -100,10 +101,22 @@ def test_login_cancelled(client):
     ],
 )
 def test_record_authentication(
-    db, sociallogin_factory, client, rf, user, process, did_record
+    db,
+    sociallogin_factory,
+    client,
+    rf,
+    user,
+    process,
+    did_record,
+    store_tokens,
+    settings,
 ):
+    settings.SOCIALACCOUNT_STORE_TOKENS = store_tokens
     sociallogin = sociallogin_factory(provider="unittest-server", uid="123")
     sociallogin.state["process"] = process
+    sociallogin.token = SocialToken(
+        app=sociallogin.account.get_provider().app, token="123", token_secret="456"
+    )
     SocialAccount.objects.create(user=user, uid="123", provider="unittest-server")
     request = rf.get("/")
     SessionMiddleware(lambda request: None).process_request(request)
@@ -122,3 +135,9 @@ def test_record_authentication(
         ]
     else:
         assert AUTHENTICATION_METHODS_SESSION_KEY not in request.session
+    assert (
+        SocialToken.objects.filter(
+            account__uid="123", token="123", token_secret="456"
+        ).exists()
+        == store_tokens
+    )