Merge pull request #349 from AzureAD/return-grouped-accounts

get_accounts() will group equivalent accounts and return the merged result

Author: rayluo

msal/application.py

@@ -642,10 +642,25 @@ def get_accounts(self, username=None):
         return accounts
 
     def _find_msal_accounts(self, environment):
-        return [a for a in self.token_cache.find(
-            TokenCache.CredentialType.ACCOUNT, query={"environment": environment})
+        grouped_accounts = {
+            a.get("home_account_id"):  # Grouped by home tenant's id
+                {  # These are minimal amount of non-tenant-specific account info
+                    "home_account_id": a.get("home_account_id"),
+                    "environment": a.get("environment"),
+                    "username": a.get("username"),
+
+                    # The following fields for backward compatibility, for now
+                    "authority_type": a.get("authority_type"),
+                    "local_account_id": a.get("local_account_id"),  # Tenant-specific
+                    "realm": a.get("realm"),  # Tenant-specific
+                }
+            for a in self.token_cache.find(
+                TokenCache.CredentialType.ACCOUNT,
+                query={"environment": environment})
             if a["authority_type"] in (
-                TokenCache.AuthorityType.ADFS, TokenCache.AuthorityType.MSSTS)]
+                TokenCache.AuthorityType.ADFS, TokenCache.AuthorityType.MSSTS)
+            }
+        return list(grouped_accounts.values())
 
     def _get_authority_aliases(self, instance):
         if not self.authority_groups:

tests/test_application.py

@@ -555,3 +555,42 @@ def mock_post(url, headers=None, *args, **kwargs):
         result = self.app.acquire_token_on_behalf_of("assertion", ["s"], post=mock_post)
         self.assertEqual(at, result.get("access_token"))
 
+
+class TestClientApplicationWillGroupAccounts(unittest.TestCase):
+    def test_get_accounts(self):
+        client_id = "my_app"
+        scopes = ["scope_1", "scope_2"]
+        environment = "login.microsoftonline.com"
+        uid = "home_oid"
+        utid = "home_tenant_guid"
+        username = "Jane Doe"
+        cache = msal.SerializableTokenCache()
+        for tenant in ["contoso", "fabrikam"]:
+            cache.add({
+                "client_id": client_id,
+                "scope": scopes,
+                "token_endpoint":
+                    "https://{}/{}/oauth2/v2.0/token".format(environment, tenant),
+                "response": TokenCacheTestCase.build_response(
+                    uid=uid, utid=utid, access_token="at", refresh_token="rt",
+                    id_token=TokenCacheTestCase.build_id_token(
+                        aud=client_id,
+                        sub="oid_in_" + tenant,
+                        preferred_username=username,
+                        ),
+                    ),
+                })
+        app = ClientApplication(
+            client_id,
+            authority="https://{}/common".format(environment),
+            token_cache=cache)
+        accounts = app.get_accounts()
+        self.assertEqual(1, len(accounts), "Should return one grouped account")
+        account = accounts[0]
+        self.assertEqual("{}.{}".format(uid, utid), account["home_account_id"])
+        self.assertEqual(environment, account["environment"])
+        self.assertEqual(username, account["username"])
+        self.assertIn("authority_type", account, "Backward compatibility")
+        self.assertIn("local_account_id", account, "Backward compatibility")
+        self.assertIn("realm", account, "Backward compatibility")
+