async def test_continuation_token():
"""Methods returning pollers should accept continuation tokens"""
expected_token = "token"
mock_generated_client = mock.Mock()
mock_methods = [
getattr(mock_generated_client, method_name) for method_name in (
"begin_full_backup",
"begin_full_restore_operation",
"begin_selective_key_restore_operation",
)
]
for method in mock_methods:
# the mock client's methods must return awaitables, and we don't have AsyncMock before 3.8
method.return_value = get_completed_future()
backup_client = KeyVaultBackupClient("vault-url", object())
backup_client._client = mock_generated_client
await backup_client.begin_restore("storage uri",
"sas",
continuation_token=expected_token)
await backup_client.begin_backup("storage uri",
"sas",
continuation_token=expected_token)
await backup_client.begin_selective_restore(
"storage uri", "sas", "key", continuation_token=expected_token)
for method in mock_methods:
assert method.call_count == 1
_, kwargs = method.call_args
assert kwargs["continuation_token"] == expected_token
async def test_full_backup_and_restore(self, container_uri, sas_token):
# backup the vault
backup_client = KeyVaultBackupClient(self.managed_hsm["url"],
self.credential)
backup_poller = await backup_client.begin_backup(
container_uri, sas_token)
# check backup status and result
job_id = backup_poller.polling_method().resource().job_id
backup_status = await backup_client.get_backup_status(job_id)
assert_in_progress_operation(backup_status)
backup_operation = await backup_poller.result()
assert_successful_operation(backup_operation)
backup_status = await backup_client.get_backup_status(job_id)
assert_successful_operation(backup_status)
# restore the backup
restore_poller = await backup_client.begin_restore(
backup_status.folder_url, sas_token)
# check restore status and result
job_id = restore_poller.polling_method().resource().job_id
restore_status = await backup_client.get_restore_status(job_id)
assert_in_progress_operation(restore_status)
restore_operation = await restore_poller.result()
assert_successful_operation(restore_operation)
restore_status = await backup_client.get_restore_status(job_id)
assert_successful_operation(restore_status)
async def test_example_backup_and_restore(self, container_uri, sas_token):
backup_client = KeyVaultBackupClient(self.managed_hsm["url"], self.credential)
# [START begin_backup]
# begin a vault backup
backup_poller = await backup_client.begin_backup(container_uri, sas_token)
# check if the backup completed
done = backup_poller.done()
# yield until the backup completes
# result() returns an object with a URL of the backup
backup_operation = await backup_poller.result()
# [END begin_backup]
folder_url = backup_operation.folder_url
# [START begin_restore]
# begin a full vault restore
restore_poller = await backup_client.begin_restore(folder_url, sas_token)
# check if the restore completed
done = restore_poller.done()
# wait for the restore to complete
await restore_poller.wait()
async def test_full_backup_and_restore_rehydration(self, container_uri,
sas_token):
if not self.is_live:
pytest.skip(
"Poller requests are incompatible with vcrpy in playback")
# backup the vault
backup_client = KeyVaultBackupClient(self.managed_hsm["url"],
self.credential)
backup_poller = await backup_client.begin_backup(
container_uri, sas_token)
# create a new poller from a continuation token
token = backup_poller.continuation_token()
rehydrated = await backup_client.begin_backup(container_uri,
sas_token,
continuation_token=token)
rehydrated_operation = await rehydrated.result()
assert rehydrated_operation.folder_url
backup_operation = await backup_poller.result()
assert backup_operation.folder_url == rehydrated_operation.folder_url
# restore the backup
restore_poller = await backup_client.begin_restore(
backup_operation.folder_url, sas_token)
# create a new poller from a continuation token
token = restore_poller.continuation_token()
rehydrated = await backup_client.begin_restore(
backup_operation.folder_url, sas_token, continuation_token=token)
await rehydrated.wait()
await restore_poller.wait()
async def test_example_selective_key_restore(self, container_uri,
sas_token):
# create a key to selectively restore
key_client = KeyClient(self.managed_hsm["url"], self.credential)
key_name = self.get_resource_name("selective-restore-test-key")
await key_client.create_rsa_key(key_name)
backup_client = KeyVaultBackupClient(self.managed_hsm["url"],
self.credential)
backup_poller = await backup_client.begin_backup(
container_uri, sas_token)
backup_operation = await backup_poller.result()
folder_url = backup_operation.folder_url
# [START begin_selective_restore]
# begin a restore of a single key from a backed up vault
restore_poller = await backup_client.begin_restore(folder_url,
sas_token,
key_name=key_name)
# check if the restore completed
done = backup_poller.done()
# wait for the restore to complete
await restore_poller.wait()
async def test_full_backup_and_restore(self, container_uri, sas_token):
# backup the vault
backup_client = KeyVaultBackupClient(self.managed_hsm["url"],
self.credential)
backup_poller = await backup_client.begin_backup(
container_uri, sas_token)
backup_operation = await backup_poller.result()
assert backup_operation.folder_url
# restore the backup
restore_poller = await backup_client.begin_restore(
backup_operation.folder_url, sas_token)
await restore_poller.wait()
async def test_full_backup_and_restore(self, container_uri, sas_token):
# backup the vault
backup_client = KeyVaultBackupClient(self.managed_hsm["url"],
self.credential)
backup_poller = await backup_client.begin_full_backup(
container_uri, sas_token)
backup_operation = await backup_poller.result()
assert_successful_operation(backup_operation)
# restore the backup
folder_name = backup_operation.azure_storage_blob_container_uri.split(
"/")[-1]
restore_poller = await backup_client.begin_full_restore(
container_uri, sas_token, folder_name)
restore_operation = await restore_poller.result()
assert_successful_operation(restore_operation)
async def test_selective_key_restore(self, container_uri, sas_token):
# create a key to selectively restore
key_client = KeyClient(self.managed_hsm["url"], self.credential)
key_name = self.get_resource_name("selective-restore-test-key")
await key_client.create_rsa_key(key_name)
# backup the vault
backup_client = KeyVaultBackupClient(self.managed_hsm["url"],
self.credential)
backup_poller = await backup_client.begin_full_backup(
container_uri, sas_token)
# check backup status and result
job_id = backup_poller.polling_method().resource().id
backup_status = await backup_client.get_backup_status(job_id)
assert_in_progress_operation(backup_status)
backup_operation = await backup_poller.result()
assert_successful_operation(backup_operation)
backup_status = await backup_client.get_backup_status(job_id)
assert_successful_operation(backup_status)
# restore the key
folder_name = backup_operation.azure_storage_blob_container_uri.split(
"/")[-1]
restore_poller = await backup_client.begin_selective_restore(
container_uri, sas_token, folder_name, key_name)
# check restore status and result
job_id = restore_poller.polling_method().resource().id
restore_status = await backup_client.get_restore_status(job_id)
assert_in_progress_operation(restore_status)
restore_operation = await restore_poller.result()
assert_successful_operation(restore_operation)
restore_status = await backup_client.get_restore_status(job_id)
assert_successful_operation(restore_status)
# delete the key
await self._poll_until_no_exception(
key_client.delete_key,
key_name,
expected_exception=ResourceExistsError)
await key_client.purge_deleted_key(key_name)
async def test_selective_key_restore(self, container_uri, sas_token):
# create a key to selectively restore
key_client = KeyClient(self.managed_hsm["url"], self.credential)
key_name = self.get_resource_name("selective-restore-test-key")
await key_client.create_rsa_key(key_name)
# backup the vault
backup_client = KeyVaultBackupClient(self.managed_hsm["url"],
self.credential)
backup_poller = await backup_client.begin_full_backup(
container_uri, sas_token)
backup_operation = await backup_poller.result()
assert_successful_operation(backup_operation)
# restore the key
folder_name = backup_operation.azure_storage_blob_container_uri.split(
"/")[-1]
restore_poller = await backup_client.begin_selective_restore(
container_uri, sas_token, folder_name, key_name)
restore_operation = await restore_poller.result()
assert_successful_operation(restore_operation)
await key_client.delete_key(key_name)
await key_client.purge_deleted_key(key_name)
async def test_selective_key_restore(self, container_uri, sas_token):
# create a key to selectively restore
key_client = KeyClient(self.managed_hsm["url"], self.credential)
key_name = self.get_resource_name("selective-restore-test-key")
await key_client.create_rsa_key(key_name)
# backup the vault
backup_client = KeyVaultBackupClient(self.managed_hsm["url"],
self.credential)
backup_poller = await backup_client.begin_backup(
container_uri, sas_token)
backup_operation = await backup_poller.result()
# restore the key
restore_poller = await backup_client.begin_restore(
backup_operation.folder_url, sas_token, key_name=key_name)
await restore_poller.wait()
# delete the key
await self._poll_until_no_exception(
key_client.delete_key,
key_name,
expected_exception=ResourceExistsError)
await key_client.purge_deleted_key(key_name)
async def test_backup_client_polling(self, container_uri, sas_token):
if not self.is_live:
pytest.skip(
"Poller requests are incompatible with vcrpy in playback")
# backup the vault
backup_client = KeyVaultBackupClient(self.managed_hsm["url"],
self.credential)
backup_poller = await backup_client.begin_backup(
container_uri, sas_token)
# create a new poller from a continuation token
token = backup_poller.continuation_token()
rehydrated = await backup_client.begin_backup(container_uri,
sas_token,
continuation_token=token)
# check that pollers and polling methods behave as expected
assert backup_poller.status() == "InProgress"
assert not backup_poller.done() or backup_poller.polling_method(
).finished()
assert rehydrated.status() == "InProgress"
assert not rehydrated.done() or rehydrated.polling_method().finished()
backup_operation = await backup_poller.result()
assert backup_poller.status(
) == "Succeeded" and backup_poller.polling_method().status(
) == "Succeeded"
rehydrated_operation = await rehydrated.result()
assert rehydrated.status(
) == "Succeeded" and rehydrated.polling_method().status(
) == "Succeeded"
assert backup_operation.folder_url == rehydrated_operation.folder_url
# rehydrate a poller with a continuation token of a completed operation
late_rehydrated = await backup_client.begin_backup(
container_uri, sas_token, continuation_token=token)
assert late_rehydrated.status() == "Succeeded"
# restore the backup
restore_poller = await backup_client.begin_restore(
backup_operation.folder_url, sas_token)
# create a new poller from a continuation token
token = restore_poller.continuation_token()
rehydrated = await backup_client.begin_restore(
backup_operation.folder_url, sas_token, continuation_token=token)
# check that pollers and polling methods behave as expected
assert restore_poller.status() == "InProgress"
assert not restore_poller.done() or restore_poller.polling_method(
).finished()
assert rehydrated.status() == "InProgress"
assert not rehydrated.done() or rehydrated.polling_method().finished()
await rehydrated.wait()
assert rehydrated.status(
) == "Succeeded" and rehydrated.polling_method().status(
) == "Succeeded"
await restore_poller.wait()
assert restore_poller.status(
) == "Succeeded" and restore_poller.polling_method().status(
) == "Succeeded"
