async def test_event_hubs_send_override_token_refresh_window(
live_eventhub_config):
uri = "sb://{}/{}".format(live_eventhub_config['hostname'],
live_eventhub_config['event_hub'])
target = "amqps://{}/{}/Partitions/0".format(
live_eventhub_config['hostname'], live_eventhub_config['event_hub'])
token = None
async def get_token():
nonlocal token
return _AccessToken(token, expiry)
jwt_auth = authentication.JWTTokenAsync(
uri,
uri,
get_token,
refresh_window=300 # set refresh window to be 5 mins
)
send_client = uamqp.SendClientAsync(target, auth=jwt_auth, debug=False)
# use token of which the valid remaining time < refresh window
expiry = int(time.time()) + (60 * 4 + 30) # 4.5 minutes
token = utils.create_sas_token(live_eventhub_config['key_name'].encode(),
live_eventhub_config['access_key'].encode(),
uri.encode(),
expiry=timedelta(minutes=4, seconds=30))
for _ in range(3):
message = uamqp.message.Message(body='Hello World')
await send_client.send_message_async(message)
auth_status = constants.CBSAuthStatus(jwt_auth._cbs_auth.get_status())
assert auth_status == constants.CBSAuthStatus.RefreshRequired
# update token, the valid remaining time > refresh window
expiry = int(time.time()) + (60 * 5 + 30) # 5.5 minutes
token = utils.create_sas_token(live_eventhub_config['key_name'].encode(),
live_eventhub_config['access_key'].encode(),
uri.encode(),
expiry=timedelta(minutes=5, seconds=30))
for _ in range(3):
message = uamqp.message.Message(body='Hello World')
await send_client.send_message_async(message)
auth_status = constants.CBSAuthStatus(jwt_auth._cbs_auth.get_status())
assert auth_status == constants.CBSAuthStatus.Ok
await send_client.close_async()
async def _create_auth_async(self) -> authentication.JWTTokenAsync:
"""
Create an ~uamqp.authentication.SASTokenAuthAsync instance to authenticate
the session.
"""
try:
# ignore mypy's warning because token_type is Optional
token_type = self._credential.token_type # type: ignore
except AttributeError:
token_type = b"jwt"
if token_type == b"servicebus.windows.net:sastoken":
auth = authentication.JWTTokenAsync(
self._auth_uri,
self._auth_uri,
functools.partial(self._credential.get_token, self._auth_uri),
token_type=token_type,
timeout=self._config.auth_timeout,
http_proxy=self._config.http_proxy,
transport_type=self._config.transport_type,
custom_endpoint_hostname=self._config.custom_endpoint_hostname,
port=self._config.connection_port,
verify=self._config.connection_verify)
await auth.update_token()
return auth
return authentication.JWTTokenAsync(
self._auth_uri,
self._auth_uri,
functools.partial(self._credential.get_token, JWT_TOKEN_SCOPE),
token_type=token_type,
timeout=self._config.auth_timeout,
http_proxy=self._config.http_proxy,
transport_type=self._config.transport_type,
custom_endpoint_hostname=self._config.custom_endpoint_hostname,
port=self._config.connection_port,
verify=self._config.connection_verify)
async def authenticate_client_by_jwt():
# Create the JWTTokenAsync object
auth_uri = "<amqp endpoint uri for authentication>" # The AMQP endpoint URI for authentication.
token_audience = "<token audience>" # The token audience field.
auth = authentication.JWTTokenAsync(audience=token_audience,
uri=auth_uri,
get_token=get_token)
# Instantiate the SendClient with the JWTTokenAsync object
target = "<target amqp service endpoint>" # The target AMQP service endpoint.
send_client = SendClientAsync(target=target, auth=auth)
# Send a message
message = Message(b'data')
await send_client.send_message_async(message)
await send_client.close_async()
def _create_auth(self, username=None, password=None):
"""
Create an ~uamqp.authentication.cbs_auth_async.SASTokenAuthAsync instance to authenticate
the session.
:param username: The name of the shared access policy.
:type username: str
:param password: The shared access key.
:type password: str
"""
http_proxy = self.config.http_proxy
transport_type = self.config.transport_type
auth_timeout = self.config.auth_timeout
if isinstance(self.credential, EventHubSharedKeyCredential):
username = username or self._auth_config['username']
password = password or self._auth_config['password']
if "@sas.root" in username:
return authentication.SASLPlain(
self.host, username, password, http_proxy=http_proxy, transport_type=transport_type)
return authentication.SASTokenAsync.from_shared_access_key(
self.auth_uri, username, password, timeout=auth_timeout, http_proxy=http_proxy,
transport_type=transport_type)
elif isinstance(self.credential, EventHubSASTokenCredential):
token = self.credential.get_sas_token()
try:
expiry = int(parse_sas_token(token)['se'])
except (KeyError, TypeError, IndexError):
raise ValueError("Supplied SAS token has no valid expiry value.")
return authentication.SASTokenAsync(
self.auth_uri, self.auth_uri, token,
expires_at=expiry,
timeout=auth_timeout,
http_proxy=http_proxy,
transport_type=transport_type)
else:
get_jwt_token = functools.partial(self.credential.get_token, 'https://eventhubs.azure.net//.default')
return authentication.JWTTokenAsync(self.auth_uri, self.auth_uri,
get_jwt_token, http_proxy=http_proxy,
transport_type=transport_type)
