|
| 1 | +"""Simulator(s) that can be used to create MSAL instance |
| 2 | +whose token cache is in a certain state, and remains unchanged. |
| 3 | +This generic simulator(s) becomes the test subject for different benchmark tools. |
| 4 | +
|
| 5 | +For example, you can install pyperf and then run: |
| 6 | +
|
| 7 | + pyperf timeit -s "from tests.simulator import ClientCredentialGrantSimulator as T; t=T(tokens_per_tenant=1, cache_hit=True)" "t.run()" |
| 8 | +""" |
| 9 | +import json |
| 10 | +import logging |
| 11 | +import random |
| 12 | +from unittest.mock import patch |
| 13 | + |
| 14 | +import msal |
| 15 | +from tests.http_client import MinimalResponse |
| 16 | + |
| 17 | + |
| 18 | +logger = logging.getLogger(__name__) |
| 19 | + |
| 20 | + |
| 21 | +def _count_access_tokens(app): |
| 22 | + return len(app.token_cache._cache[app.token_cache.CredentialType.ACCESS_TOKEN]) |
| 23 | + |
| 24 | + |
| 25 | +class ClientCredentialGrantSimulator(object): |
| 26 | + |
| 27 | + def __init__(self, number_of_tenants=1, tokens_per_tenant=1, cache_hit=False): |
| 28 | + logger.info( |
| 29 | + "number_of_tenants=%d, tokens_per_tenant=%d, cache_hit=%s", |
| 30 | + number_of_tenants, tokens_per_tenant, cache_hit) |
| 31 | + with patch.object(msal.authority, "tenant_discovery", return_value={ |
| 32 | + "authorization_endpoint": "https://contoso.com/placeholder", |
| 33 | + "token_endpoint": "https://contoso.com/placeholder", |
| 34 | + }) as _: # Otherwise it would fail on OIDC discovery |
| 35 | + self.apps = [ # In MSAL Python, each CCA binds to one tenant only |
| 36 | + msal.ConfidentialClientApplication( |
| 37 | + "client_id", client_credential="foo", |
| 38 | + authority="https://login.microsoftonline.com/tenant_%s" % t, |
| 39 | + ) for t in range(number_of_tenants) |
| 40 | + ] |
| 41 | + for app in self.apps: |
| 42 | + for i in range(tokens_per_tenant): # Populate token cache |
| 43 | + self.run(app=app, scope="scope_{}".format(i)) |
| 44 | + assert tokens_per_tenant == _count_access_tokens(app), ( |
| 45 | + "Token cache did not populate correctly: {}".format(json.dumps( |
| 46 | + app.token_cache._cache, indent=4))) |
| 47 | + |
| 48 | + if cache_hit: |
| 49 | + self.run(app=app)["access_token"] # Populate 1 token to be hit |
| 50 | + expected_tokens = tokens_per_tenant + 1 |
| 51 | + else: |
| 52 | + expected_tokens = tokens_per_tenant |
| 53 | + app.token_cache.modify = lambda *args, **kwargs: None # Cache becomes read-only |
| 54 | + self.run(app=app)["access_token"] |
| 55 | + assert expected_tokens == _count_access_tokens(app), "Cache shall not grow" |
| 56 | + |
| 57 | + def run(self, app=None, scope=None): |
| 58 | + # This implementation shall be as concise as possible |
| 59 | + app = app or random.choice(self.apps) |
| 60 | + #return app.acquire_token_for_client([scope or "scope"], post=_fake) |
| 61 | + return app.acquire_token_for_client( |
| 62 | + [scope or "scope"], |
| 63 | + post=lambda url, **kwargs: MinimalResponse( # Using an inline lambda is as fast as a standalone function |
| 64 | + status_code=200, text='''{ |
| 65 | + "access_token": "AT for %s", |
| 66 | + "token_type": "bearer" |
| 67 | + }''' % kwargs["data"]["scope"], |
| 68 | + )) |
| 69 | + |
0 commit comments