feat: always match source account in restore and volume clone scenarios

kubernetes-sigs · Nov 28, 2024 · a23c548 · a23c548
1 parent 793451a
commit a23c548
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 3 deletions.
diff --git a/pkg/provider/azure_storageaccount.go b/pkg/provider/azure_storageaccount.go
@@ -87,6 +87,8 @@ type AccountOptions struct {
 	SoftDeleteContainers                    int32
 	// indicate whether to get a random matching account, if false, will get the first matching account
 	PickRandomMatchingAccount bool
+	// provide the source account name in snapshot restore and volume clone scenarios
+	SourceAccountName string
 }
 
 type accountWithLocation struct {
@@ -341,6 +343,15 @@ func (az *Cloud) EnsureStorageAccount(ctx context.Context, accountOptions *Accou
 				}
 				accountName = accounts[index].Name
 				createNewAccount = false
+				if accountOptions.SourceAccountName != "" {
+					for _, acct := range accounts {
+						if acct.Name == accountOptions.SourceAccountName {
+							klog.V(2).Infof("found a matching account %s type %s location %s with source account name", acct.Name, acct.StorageType, acct.Location)
+							accountName = acct.Name
+							break
+						}
+					}
+				}
 				klog.V(4).Infof("found a matching account %s type %s location %s", accounts[index].Name, accounts[index].StorageType, accounts[index].Location)
 			}
 		}

diff --git a/pkg/provider/azure_storageaccount_test.go b/pkg/provider/azure_storageaccount_test.go
@@ -395,14 +395,16 @@ func TestEnsureStorageAccount(t *testing.T) {
 	cloud.Location = location
 	cloud.SubscriptionID = "testSub"
 
-	name := "testStorageAccount"
 	sku := &storage.Sku{
 		Name: "testSku",
 		Tier: "testSkuTier",
 	}
 	testStorageAccounts :=
 		[]storage.Account{
-			{Name: &name, Kind: "kind", Location: &location, Sku: sku, AccountProperties: &storage.AccountProperties{NetworkRuleSet: &storage.NetworkRuleSet{}}}}
+			{Name: ptr.To("testStorageAccount"), Kind: "kind", Location: &location, Sku: sku, AccountProperties: &storage.AccountProperties{NetworkRuleSet: &storage.NetworkRuleSet{}}},
+			{Name: ptr.To("wantedAccount"), Kind: "kind", Location: &location, Sku: sku, AccountProperties: &storage.AccountProperties{NetworkRuleSet: &storage.NetworkRuleSet{}}},
+			{Name: ptr.To("otherAccount"), Kind: "kind", Location: &location, Sku: sku, AccountProperties: &storage.AccountProperties{NetworkRuleSet: &storage.NetworkRuleSet{}}},
+		}
 
 	value := "foo bar"
 	storageAccountListKeys := storage.AccountListKeysResult{
@@ -423,10 +425,12 @@ func TestEnsureStorageAccount(t *testing.T) {
 		storageType                     StorageType
 		requireInfrastructureEncryption *bool
 		keyVaultURL                     *string
+		sourceAccountName               string
 		accountName                     string
 		subscriptionID                  string
 		resourceGroup                   string
 		expectedErr                     string
+		expectedAccountName             string
 	}{
 		{
 			name:                            "[Success] EnsureStorageAccount with createPrivateEndpoint and storagetype blob",
@@ -456,6 +460,16 @@ func TestEnsureStorageAccount(t *testing.T) {
 			accountName:                     "",
 			expectedErr:                     "",
 		},
+		{
+			name:                            "[Success] EnsureStorageAccount returns with source account",
+			mockStorageAccountsClient:       true,
+			setAccountOptions:               true,
+			requireInfrastructureEncryption: ptr.To(true),
+			resourceGroup:                   "rg",
+			sourceAccountName:               "wantedAccount",
+			accountName:                     "wantedAccount",
+			expectedErr:                     "",
+		},
 		{
 			name:                      "[Failed] EnsureStorageAccount with createPrivateEndpoint: get storage key failed",
 			createAccount:             true,
@@ -537,6 +551,10 @@ func TestEnsureStorageAccount(t *testing.T) {
 			mockVirtualNetworkLinksClient.EXPECT().CreateOrUpdate(gomock.Any(), vnetResourceGroup, gomock.Any(), gomock.Any(), gomock.Any()).Return(nil, nil).Times(1)
 		}
 
+		if test.sourceAccountName != "" {
+			mockStorageAccountsClient.EXPECT().ListKeys(gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any()).Return(storageAccountListKeys, nil).AnyTimes()
+		}
+
 		var testAccountOptions *AccountOptions
 		if test.setAccountOptions {
 			testAccountOptions = &AccountOptions{
@@ -551,10 +569,14 @@ func TestEnsureStorageAccount(t *testing.T) {
 				SoftDeleteBlobs:           7,
 				SoftDeleteContainers:      7,
 				PickRandomMatchingAccount: test.pickRandomMatchingAccount,
+				SourceAccountName:         test.sourceAccountName,
 			}
 		}
 
-		_, _, err := cloud.EnsureStorageAccount(ctx, testAccountOptions, "test")
+		accountName, _, err := cloud.EnsureStorageAccount(ctx, testAccountOptions, "test")
+		if test.expectedAccountName != "" {
+			assert.Equal(t, accountName, test.expectedAccountName, test.name)
+		}
 		assert.Equal(t, err == nil, test.expectedErr == "", fmt.Sprintf("returned error: %v", err), test.name)
 		if test.expectedErr != "" {
 			assert.Equal(t, err != nil, strings.Contains(err.Error(), test.expectedErr), err.Error(), test.name)