Add StateChangeConf

hashicorp · Jun 22, 2021 · cb4841a · cb4841a
1 parent b7653cd
commit cb4841a
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 5 deletions.
diff --git a/...m/internal/services/dataprotection/data_protection_backup_instance_postgresql_resource.go b/...m/internal/services/dataprotection/data_protection_backup_instance_postgresql_resource.go
@@ -1,6 +1,7 @@
 package dataprotection
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"time"
@@ -16,6 +17,7 @@ import (
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/services/dataprotection/validate"
 	postgresParse "github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/services/postgres/parse"
 	postgresValidate "github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/services/postgres/validate"
+	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/tf/pluginsdk"
 	azSchema "github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/tf/schema"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/internal/timeouts"
 	"github.com/terraform-providers/terraform-provider-azurerm/azurerm/utils"
@@ -136,6 +138,18 @@ func resourceDataProtectionBackupInstancePostgreSQLCreateUpdate(d *schema.Resour
 		return fmt.Errorf("waiting for creation/update of the DataProtection BackupInstance (%q): %+v", id, err)
 	}
 
+	stateConf := &pluginsdk.StateChangeConf{
+		Pending:    []string{string(dataprotection.ConfiguringProtection), string(dataprotection.UpdatingProtection)},
+		Target:     []string{string(dataprotection.ProtectionConfigured)},
+		Refresh:    policyProtectionStateRefreshFunc(ctx, client, id),
+		MinTimeout: 1 * time.Minute,
+		Timeout:    d.Timeout(pluginsdk.TimeoutCreate),
+	}
+
+	if _, err = stateConf.WaitForState(); err != nil {
+		return fmt.Errorf("waiting for BackupInstance(%q) policy protection to be completed: %+v", id, err)
+	}
+
 	d.SetId(id.ID())
 	return resourceDataProtectionBackupInstancePostgreSQLRead(d, meta)
 }
@@ -197,3 +211,14 @@ func resourceDataProtectionBackupInstancePostgreSQLDelete(d *schema.ResourceData
 	}
 	return nil
 }
+
+func policyProtectionStateRefreshFunc(ctx context.Context, client *dataprotection.BackupInstancesClient, id parse.BackupInstanceId) pluginsdk.StateRefreshFunc {
+	return func() (interface{}, string, error) {
+		res, err := client.Get(ctx, id.BackupVaultName, id.ResourceGroup, id.Name)
+		if err != nil {
+			return nil, "", fmt.Errorf("retrieving DataProtection BackupInstance (%q): %+v", id, err)
+		}
+
+		return res, string(res.Properties.ProtectionStatus.Status), nil
+	}
+}
diff --git a/...ernal/services/dataprotection/data_protection_backup_instance_postgresql_resource_test.go b/...ernal/services/dataprotection/data_protection_backup_instance_postgresql_resource_test.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"fmt"
 	"testing"
-	"time"
 
 	"github.com/hashicorp/terraform-plugin-sdk/helper/resource"
 	"github.com/hashicorp/terraform-plugin-sdk/terraform"
@@ -71,16 +70,14 @@ func TestAccDataProtectionBackupInstancePostgreSQL_update(t *testing.T) {
 		},
 		data.ImportStep(),
 		{
-			PreConfig: func() { time.Sleep(3 * time.Minute) },
-			Config:    r.complete(data),
+			Config: r.complete(data),
 			Check: resource.ComposeTestCheckFunc(
 				check.That(data.ResourceName).ExistsInAzure(r),
 			),
 		},
 		data.ImportStep(),
 		{
-			PreConfig: func() { time.Sleep(3 * time.Minute) },
-			Config:    r.basic(data),
+			Config: r.basic(data),
 			Check: resource.ComposeTestCheckFunc(
 				check.That(data.ResourceName).ExistsInAzure(r),
 			),