def test_imds_credential_retries():
mock_response = Mock(
text=lambda: b"{}",
headers={"content-type": "application/json", "Retry-After": "0"},
content_type="application/json",
)
mock_send = Mock(return_value=mock_response)
total_retries = ImdsCredential._create_config().retry_policy.total_retries
for status_code in (404, 429, 500):
mock_send.reset_mock()
mock_response.status_code = status_code
try:
ImdsCredential(transport=Mock(send=mock_send)).get_token("scope")
except ClientAuthenticationError:
pass
# first call was availability probe, second the original request;
# credential should have then exhausted retries for each of these status codes
assert mock_send.call_count == 2 + total_retries
def test_imds_credential_cache():
scope = "https://foo.bar"
expired = "this token's expired"
now = int(time.time())
token_payload = {
"access_token": expired,
"refresh_token": "",
"expires_in": 0,
"expires_on": now - 300, # expired 5 minutes ago
"not_before": now,
"resource": scope,
"token_type": "Bearer",
}
mock_response = Mock(
text=lambda: json.dumps(token_payload),
headers={"content-type": "application/json"},
status_code=200,
content_type="application/json",
)
mock_send = Mock(return_value=mock_response)
credential = ImdsCredential(transport=Mock(send=mock_send))
token = credential.get_token(scope)
assert token.token == expired
assert mock_send.call_count == 2 # first request was probing for endpoint availability
# calling get_token again should provoke another HTTP request
good_for_an_hour = "this token's good for an hour"
token_payload["expires_on"] = int(time.time()) + 3600
token_payload["expires_in"] = 3600
token_payload["access_token"] = good_for_an_hour
token = credential.get_token(scope)
assert token.token == good_for_an_hour
assert mock_send.call_count == 3
# get_token should return the cached token now
token = credential.get_token(scope)
assert token.token == good_for_an_hour
assert mock_send.call_count == 3