From ae602330f655453f2ae18557715324bc2bcec12f Mon Sep 17 00:00:00 2001
From: Grace Wilcox <43627800+gracewilcox@users.noreply.github.com>
Date: Wed, 25 Sep 2024 14:26:02 -0700
Subject: [PATCH] [azadmin] change frequency of poller in tests (#23475)

---
 .../keyvault/azadmin/backup/client_test.go    | 41 +++++++++++++++----
 1 file changed, 34 insertions(+), 7 deletions(-)

diff --git a/sdk/security/keyvault/azadmin/backup/client_test.go b/sdk/security/keyvault/azadmin/backup/client_test.go
index addfaaf084dd..bbb06509562c 100644
--- a/sdk/security/keyvault/azadmin/backup/client_test.go
+++ b/sdk/security/keyvault/azadmin/backup/client_test.go
@@ -12,6 +12,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/Azure/azure-sdk-for-go/sdk/azcore/runtime"
 	"github.com/Azure/azure-sdk-for-go/sdk/azcore/to"
 	"github.com/Azure/azure-sdk-for-go/sdk/internal/recording"
 	"github.com/Azure/azure-sdk-for-go/sdk/security/keyvault/azadmin/backup"
@@ -21,13 +22,19 @@ import (
 
 func TestBackupRestore(t *testing.T) {
 	client, sasToken := startBackupTest(t)
-
 	testSerde(t, &sasToken)
 
+	var frequency time.Duration
+	if recording.GetRecordMode() == recording.PlaybackMode {
+		frequency = time.Second
+	}
+
 	// backup the vault
 	backupPoller, err := client.BeginFullBackup(context.Background(), sasToken, nil)
 	require.NoError(t, err)
-	backupResults, err := backupPoller.PollUntilDone(context.Background(), nil)
+	backupResults, err := backupPoller.PollUntilDone(context.Background(), &runtime.PollUntilDoneOptions{
+		Frequency: frequency,
+	})
 	require.NoError(t, err)
 	require.Nil(t, backupResults.Error)
 	require.Equal(t, "Succeeded", *backupResults.Status)
@@ -44,7 +51,9 @@ func TestBackupRestore(t *testing.T) {
 	testSerde(t, &restoreOperationParameters)
 	restorePoller, err := client.BeginFullRestore(context.Background(), restoreOperationParameters, nil)
 	require.NoError(t, err)
-	restoreResults, err := restorePoller.PollUntilDone(context.Background(), nil)
+	restoreResults, err := restorePoller.PollUntilDone(context.Background(), &runtime.PollUntilDoneOptions{
+		Frequency: frequency,
+	})
 	require.NoError(t, err)
 	require.Nil(t, restoreResults.Error)
 	require.Equal(t, "Succeeded", *restoreResults.Status)
@@ -62,6 +71,11 @@ func TestBackupRestore(t *testing.T) {
 func TestBackupRestoreWithResumeToken(t *testing.T) {
 	client, sasToken := startBackupTest(t)
 
+	var frequency time.Duration
+	if recording.GetRecordMode() == recording.PlaybackMode {
+		frequency = time.Second
+	}
+
 	// backup the vault
 	backupPoller, err := client.BeginFullBackup(context.Background(), sasToken, nil)
 	require.NoError(t, err)
@@ -71,7 +85,9 @@ func TestBackupRestoreWithResumeToken(t *testing.T) {
 	require.NoError(t, err)
 	newBackupPoller, err := client.BeginFullBackup(context.Background(), sasToken, &backup.BeginFullBackupOptions{ResumeToken: token})
 	require.NoError(t, err)
-	backupResults, err := newBackupPoller.PollUntilDone(context.Background(), nil)
+	backupResults, err := newBackupPoller.PollUntilDone(context.Background(), &runtime.PollUntilDoneOptions{
+		Frequency: frequency,
+	})
 	require.NoError(t, err)
 	require.Nil(t, backupResults.Error)
 	require.Equal(t, "Succeeded", *backupResults.Status)
@@ -94,7 +110,9 @@ func TestBackupRestoreWithResumeToken(t *testing.T) {
 	require.NoError(t, err)
 	newRestorePoller, err := client.BeginFullRestore(context.Background(), restoreOperationParameters, &backup.BeginFullRestoreOptions{ResumeToken: restoreToken})
 	require.NoError(t, err)
-	restoreResults, err := newRestorePoller.PollUntilDone(context.Background(), nil)
+	restoreResults, err := newRestorePoller.PollUntilDone(context.Background(), &runtime.PollUntilDoneOptions{
+		Frequency: frequency,
+	})
 	require.NoError(t, err)
 	require.Nil(t, restoreResults.Error)
 	require.Equal(t, "Succeeded", *restoreResults.Status)
@@ -112,6 +130,11 @@ func TestBackupRestoreWithResumeToken(t *testing.T) {
 func TestBeginSelectiveKeyRestoreOperation(t *testing.T) {
 	backupClient, sasToken := startBackupTest(t)
 
+	var frequency time.Duration
+	if recording.GetRecordMode() == recording.PlaybackMode {
+		frequency = time.Second
+	}
+
 	// create a key to selectively restore
 	if recording.GetRecordMode() != recording.PlaybackMode {
 		cred := credential
@@ -128,7 +151,9 @@ func TestBeginSelectiveKeyRestoreOperation(t *testing.T) {
 	// backup the vault
 	backupPoller, err := backupClient.BeginFullBackup(context.Background(), sasToken, nil)
 	require.NoError(t, err)
-	backupResults, err := backupPoller.PollUntilDone(context.Background(), nil)
+	backupResults, err := backupPoller.PollUntilDone(context.Background(), &runtime.PollUntilDoneOptions{
+		Frequency: frequency,
+	})
 	require.NoError(t, err)
 
 	// restore the key
@@ -141,7 +166,9 @@ func TestBeginSelectiveKeyRestoreOperation(t *testing.T) {
 	testSerde(t, &restoreOperationParameters)
 	selectivePoller, err := backupClient.BeginSelectiveKeyRestore(context.Background(), "selective-restore-test-key", restoreOperationParameters, nil)
 	require.NoError(t, err)
-	selectiveResults, err := selectivePoller.PollUntilDone(context.Background(), nil)
+	selectiveResults, err := selectivePoller.PollUntilDone(context.Background(), &runtime.PollUntilDoneOptions{
+		Frequency: frequency,
+	})
 	require.NoError(t, err)
 	testSerde(t, &selectiveResults)