From 66e7456d0075e2737be8b702e90ab61d1fd872d6 Mon Sep 17 00:00:00 2001
From: Owen Farrell <owen.farrell@gmail.com>
Date: Mon, 27 Sep 2021 17:00:46 -0400
Subject: [PATCH] Wait for firewall rule provisioning state when creating or
 updating resource

Signed-off-by: Owen Farrell <owen.farrell@gmail.com>
---
 .../synapse/synapse_firewall_rule_resource.go | 30 +++++++++++++++++--
 1 file changed, 27 insertions(+), 3 deletions(-)

diff --git a/internal/services/synapse/synapse_firewall_rule_resource.go b/internal/services/synapse/synapse_firewall_rule_resource.go
index 672ef5372d661..74b5bd24ec9d4 100644
--- a/internal/services/synapse/synapse_firewall_rule_resource.go
+++ b/internal/services/synapse/synapse_firewall_rule_resource.go
@@ -1,6 +1,7 @@
 package synapse
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"time"
@@ -105,9 +106,17 @@ func resourceSynapseFirewallRuleCreateUpdate(d *pluginsdk.ResourceData, meta int
 		return fmt.Errorf("waiting for creation/update of %s: %+v", id, err)
 	}
 
-	_, err = client.Get(ctx, id.ResourceGroup, id.WorkspaceName, id.Name)
-	if err != nil {
-		return fmt.Errorf("retrieving %s: %+v", id, err)
+	timeout, _ := ctx.Deadline()
+
+	stateConf := &pluginsdk.StateChangeConf{
+		Pending:    []string{string(synapse.ProvisioningStateProvisioning)},
+		Target:     []string{string(synapse.ProvisioningStateSucceeded)},
+		Refresh:    firewallRuleProvisioningStateRefreshFunc(ctx, client, id),
+		MinTimeout: 1 * time.Minute,
+		Timeout:    time.Until(timeout),
+	}
+	if _, err = stateConf.WaitForStateContext(ctx); err != nil {
+		return fmt.Errorf("waiting for provisioning state of %s: %+v", id, err)
 	}
 
 	d.SetId(id.ID())
@@ -169,3 +178,18 @@ func resourceSynapseFirewallRuleDelete(d *pluginsdk.ResourceData, meta interface
 
 	return nil
 }
+
+func firewallRuleProvisioningStateRefreshFunc(ctx context.Context, client *synapse.IPFirewallRulesClient, id parse.FirewallRuleId) pluginsdk.StateRefreshFunc {
+	return func() (interface{}, string, error) {
+		resp, err := client.Get(ctx, id.ResourceGroup, id.WorkspaceName, id.Name)
+		if err != nil {
+			return nil, "", fmt.Errorf("polling for %s: %+v", id, err)
+		}
+
+		if resp.IPFirewallRuleProperties == nil {
+			return resp, "", fmt.Errorf("nil IPFirewallRuleProperties returned for %s", id)
+		}
+
+		return resp, string(resp.IPFirewallRuleProperties.ProvisioningState), nil
+	}
+}