Reuse old rt data even if its key is different

Author: rayluo

msal/token_cache.py

@@ -234,8 +234,9 @@ def modify(self, credential_type, old_entry, new_key_value_pairs=None):
         with self._lock:
             if new_key_value_pairs:  # Update with them
                 entries = self._cache.setdefault(credential_type, {})
-                entry = entries.setdefault(key, {})  # Create it if not yet exist
-                entry.update(new_key_value_pairs)
+                entries[key] = dict(
+                    old_entry,  # Do not use entries[key] b/c it might not exist
+                    **new_key_value_pairs)
             else:  # Remove old_entry
                 self._cache.setdefault(credential_type, {}).pop(key, None)
 

tests/test_token_cache.py

@@ -222,6 +222,24 @@ def test_key_id_is_also_recorded(self):
             {}).get("key_id")
         self.assertEqual(my_key_id, cached_key_id, "AT should be bound to the key")
 
+    def test_old_rt_data_with_wrong_key_should_still_be_salvaged_into_new_rt(self):
+        sample = {
+            'client_id': 'my_client_id',
+            'credential_type': 'RefreshToken',
+            'environment': 'login.example.com',
+            'home_account_id': "uid.utid",
+            'secret': 'a refresh token',
+            'target': 's2 s1 s3',
+            }
+        new_rt = "this is a new RT"
+        self.cache._cache["RefreshToken"] = {"wrong-key": sample}
+        self.cache.modify(
+            self.cache.CredentialType.REFRESH_TOKEN, sample, {"secret": new_rt})
+        self.assertEqual(
+            dict(sample, secret=new_rt),
+            self.cache._cache["RefreshToken"].get(
+                'uid.utid-login.example.com-refreshtoken-my_client_id--s2 s1 s3')
+            )
 
 class SerializableTokenCacheTestCase(TokenCacheTestCase):
     # Run all inherited test methods, and have extra check in tearDown()