From 78db139da93377fd2f332b5625d460bf3ce8efc7 Mon Sep 17 00:00:00 2001
From: magodo <wztdyl@sina.com>
Date: Fri, 22 Sep 2023 14:38:52 +0800
Subject: [PATCH] `azurerm_management_lock` - Add post-create/delete polling to
 tolerate RP propagation (#23345)

* `azurerm_management_lock` - Add post-create/delete polling to tolerate RP propagation

* Using ctx deadline
---
 .../resource/management_lock_resource.go      | 51 ++++++++++++++++++-
 1 file changed, 50 insertions(+), 1 deletion(-)

diff --git a/internal/services/resource/management_lock_resource.go b/internal/services/resource/management_lock_resource.go
index c083a0f08995..08ea57aa062f 100644
--- a/internal/services/resource/management_lock_resource.go
+++ b/internal/services/resource/management_lock_resource.go
@@ -4,6 +4,7 @@
 package resource
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -32,7 +33,6 @@ func resourceManagementLock() *pluginsdk.Resource {
 		Timeouts: &pluginsdk.ResourceTimeout{
 			Create: pluginsdk.DefaultTimeout(30 * time.Minute),
 			Read:   pluginsdk.DefaultTimeout(5 * time.Minute),
-			Update: pluginsdk.DefaultTimeout(30 * time.Minute),
 			Delete: pluginsdk.DefaultTimeout(30 * time.Minute),
 		},
 
@@ -100,6 +100,24 @@ func resourceManagementLockCreate(d *pluginsdk.ResourceData, meta interface{}) e
 		return fmt.Errorf("creating %s: %+v", id, err)
 	}
 
+	deadline, ok := ctx.Deadline()
+	if !ok {
+		return fmt.Errorf("internal-error: context was missing a deadline")
+	}
+
+	stateConf := &pluginsdk.StateChangeConf{
+		Target: []string{
+			"OK",
+		},
+		Refresh:                   managementLockStateRefreshFunc(ctx, client, id),
+		MinTimeout:                10 * time.Second,
+		ContinuousTargetOccurence: 12,
+		Timeout:                   time.Until(deadline),
+	}
+	if _, err := stateConf.WaitForStateContext(ctx); err != nil {
+		return fmt.Errorf("waiting for %s to finish create replication", id)
+	}
+
 	d.SetId(id.ID())
 	return resourceManagementLockRead(d, meta)
 }
@@ -155,5 +173,36 @@ func resourceManagementLockDelete(d *pluginsdk.ResourceData, meta interface{}) e
 		return fmt.Errorf("deleting %s: %+v", *id, err)
 	}
 
+	deadline, ok := ctx.Deadline()
+	if !ok {
+		return fmt.Errorf("internal-error: context was missing a deadline")
+	}
+
+	stateConf := &pluginsdk.StateChangeConf{
+		Target: []string{
+			"NotFound",
+		},
+		Refresh:                   managementLockStateRefreshFunc(ctx, client, *id),
+		MinTimeout:                10 * time.Second,
+		ContinuousTargetOccurence: 12,
+		Timeout:                   time.Until(deadline),
+	}
+	if _, err := stateConf.WaitForStateContext(ctx); err != nil {
+		return fmt.Errorf("waiting for %s to finish delete replication", id)
+	}
+
 	return nil
 }
+
+func managementLockStateRefreshFunc(ctx context.Context, client *managementlocks.ManagementLocksClient, id managementlocks.ScopedLockId) pluginsdk.StateRefreshFunc {
+	return func() (interface{}, string, error) {
+		resp, err := client.GetByScope(ctx, id)
+		if err != nil {
+			if response.WasNotFound(resp.HttpResponse) {
+				return resp, "NotFound", nil
+			}
+			return nil, "Error", err
+		}
+		return "OK", "OK", nil
+	}
+}