Merge pull request #247 from AzureAD/cache-tokens-by-specified-environment

Cache tokens by specified environment, not by OIDC Discovery

Author: rayluo

msal/application.py

@@ -274,7 +274,8 @@ def _build_client(self, client_credential, authority):
             default_body=default_body,
             client_assertion=client_assertion,
             client_assertion_type=client_assertion_type,
-            on_obtaining_tokens=self.token_cache.add,
+            on_obtaining_tokens=lambda event: self.token_cache.add(dict(
+                event, environment=authority.instance)),
             on_removing_rt=self.token_cache.remove_rt,
             on_updating_rt=self.token_cache.update_rt)
 

msal/token_cache.py

@@ -126,6 +126,8 @@ def __add(self, event, now=None):
         environment = realm = None
         if "token_endpoint" in event:
             _, environment, realm = canonicalize(event["token_endpoint"])
+        if "environment" in event:  # Always available unless in legacy test cases
+            environment = event["environment"]  # Set by application.py
         response = event.get("response", {})
         data = event.get("data", {})
         access_token = response.get("access_token")