def test_continuation_token():
"""Methods returning pollers should accept continuation tokens"""
expected_token = "token"
mock_generated_client = mock.Mock()
backup_client = KeyVaultBackupClient("vault-url", object())
backup_client._client = mock_generated_client
backup_client.begin_full_restore("storage uri",
"sas",
"folder",
continuation_token=expected_token)
backup_client.begin_full_backup("storage uri",
"sas",
continuation_token=expected_token)
backup_client.begin_selective_restore("storage uri",
"sas",
"folder",
"key",
continuation_token=expected_token)
for method in ("begin_full_backup", "begin_full_restore_operation",
"begin_selective_key_restore_operation"):
mock_method = getattr(mock_generated_client, method)
assert mock_method.call_count == 1
_, kwargs = mock_method.call_args
assert kwargs["continuation_token"] == expected_token
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")
key_client.create_rsa_key(key_name)
# backup the vault
backup_client = KeyVaultBackupClient(self.managed_hsm["url"],
self.credential)
backup_poller = backup_client.begin_full_backup(
container_uri, sas_token)
# check backup status and result
job_id = backup_poller.polling_method().resource().id
backup_status = backup_client.get_backup_status(job_id)
assert_in_progress_operation(backup_status)
backup_operation = backup_poller.result()
assert_successful_operation(backup_operation)
backup_status = 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 = 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 = backup_client.get_restore_status(job_id)
assert_in_progress_operation(restore_status)
restore_operation = restore_poller.result()
assert_successful_operation(restore_operation)
restore_status = backup_client.get_restore_status(job_id)
assert_successful_operation(restore_status)
# delete the key
delete_function = partial(key_client.begin_delete_key, key_name)
delete_poller = self._poll_until_no_exception(delete_function,
ResourceExistsError)
delete_poller.wait()
key_client.purge_deleted_key(key_name)
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")
key_client.create_rsa_key(key_name)
# backup the vault
backup_client = KeyVaultBackupClient(self.managed_hsm["url"], self.credential)
backup_poller = backup_client.begin_full_backup(container_uri, sas_token)
backup_operation = backup_poller.result()
assert_successful_operation(backup_operation)
# restore the key
folder_name = backup_operation.azure_storage_blob_container_uri.split("/")[-1]
restore_poller = backup_client.begin_selective_restore(container_uri, sas_token, folder_name, key_name)
restore_operation = restore_poller.result()
assert_successful_operation(restore_operation)
key_client.begin_delete_key(key_name).wait()
key_client.purge_deleted_key(key_name)