Prevent crash on token_cache.find(..., query=None)

AzureAD · Jan 9, 2024 · 84bdfab · 84bdfab
1 parent 5272fbd
commit 84bdfab
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/msal/token_cache.py b/msal/token_cache.py
@@ -128,6 +128,7 @@ def _find(self, credential_type, target=None, query=None):  # O(n) generator
 
         preferred_result = None
         if (credential_type == self.CredentialType.ACCESS_TOKEN
+            and isinstance(query, dict)
             and "home_account_id" in query and "environment" in query
             and "client_id" in query and "realm" in query and target
         ):  # Special case for O(1) AT lookup

diff --git a/tests/test_token_cache.py b/tests/test_token_cache.py
@@ -65,8 +65,7 @@ def testAddByAad(self):
                 expires_in=3600, access_token="an access token",
                 id_token=id_token, refresh_token="a refresh token"),
             }, now=1000)
-        self.assertEqual(
-            {
+        access_token_entry = {
                 'cached_at': "1000",
                 'client_id': 'my_client_id',
                 'credential_type': 'AccessToken',
@@ -78,10 +77,16 @@ def testAddByAad(self):
                 'secret': 'an access token',
                 'target': 's1 s2 s3',  # Sorted
                 'token_type': 'some type',
-            },
+            }
+        self.assertEqual(
+            access_token_entry,
             self.cache._cache["AccessToken"].get(
                 'uid.utid-login.example.com-accesstoken-my_client_id-contoso-s1 s2 s3')
             )
+        self.assertIn(
+            access_token_entry,
+            self.cache.find(self.cache.CredentialType.ACCESS_TOKEN),
+            "find(..., query=None) should not crash, even though MSAL does not use it")
         self.assertEqual(
             {
                 'client_id': 'my_client_id',