From 15921b02bad313d981d34d7862ee95ee5616740a Mon Sep 17 00:00:00 2001
From: Yichun Ma <yicma@microsoft.com>
Date: Sun, 28 Aug 2022 23:29:31 +0800
Subject: [PATCH] `azurerm_iothub_dps` - support `data_residency_enabled`

---
 .../services/iothub/iothub_dps_resource.go    | 15 +++++++
 .../iothub/iothub_dps_resource_test.go        | 44 +++++++++++++++++++
 website/docs/r/iothub_dps.html.markdown       |  2 +
 3 files changed, 61 insertions(+)

diff --git a/internal/services/iothub/iothub_dps_resource.go b/internal/services/iothub/iothub_dps_resource.go
index 834809f43482..77ae4f7408dc 100644
--- a/internal/services/iothub/iothub_dps_resource.go
+++ b/internal/services/iothub/iothub_dps_resource.go
@@ -170,6 +170,13 @@ func resourceIotHubDPS() *pluginsdk.Resource {
 				}, false),
 			},
 
+			"data_residency_enabled": {
+				Type:     pluginsdk.TypeBool,
+				Optional: true,
+				ForceNew: true,
+				Default:  false,
+			},
+
 			"public_network_access_enabled": {
 				Type:     pluginsdk.TypeBool,
 				Optional: true,
@@ -229,6 +236,7 @@ func resourceIotHubDPSCreateUpdate(d *pluginsdk.ResourceData, meta interface{})
 		Properties: &iothub.IotDpsPropertiesDescription{
 			IotHubs:             expandIoTHubDPSIoTHubs(d.Get("linked_hub").([]interface{})),
 			AllocationPolicy:    iothub.AllocationPolicy(d.Get("allocation_policy").(string)),
+			EnableDataResidency: utils.Bool(d.Get("data_residency_enabled").(bool)),
 			IPFilterRules:       expandDpsIPFilterRules(d),
 			PublicNetworkAccess: publicNetworkAccess,
 		},
@@ -293,6 +301,13 @@ func resourceIotHubDPSRead(d *pluginsdk.ResourceData, meta interface{}) error {
 		d.Set("device_provisioning_host_name", props.DeviceProvisioningHostName)
 		d.Set("id_scope", props.IDScope)
 		d.Set("allocation_policy", string(props.AllocationPolicy))
+
+		enableDataResidency := false
+		if props.EnableDataResidency != nil {
+			enableDataResidency = *props.EnableDataResidency
+		}
+		d.Set("data_residency_enabled", enableDataResidency)
+
 		publicNetworkAccess := true
 		if props.PublicNetworkAccess != "" {
 			publicNetworkAccess = strings.EqualFold("Enabled", string(props.PublicNetworkAccess))
diff --git a/internal/services/iothub/iothub_dps_resource_test.go b/internal/services/iothub/iothub_dps_resource_test.go
index 875c85d154fe..40e43881c1a2 100644
--- a/internal/services/iothub/iothub_dps_resource_test.go
+++ b/internal/services/iothub/iothub_dps_resource_test.go
@@ -28,6 +28,7 @@ func TestAccIotHubDPS_basic(t *testing.T) {
 				check.That(data.ResourceName).Key("device_provisioning_host_name").Exists(),
 				check.That(data.ResourceName).Key("id_scope").Exists(),
 				check.That(data.ResourceName).Key("service_operations_host_name").Exists(),
+				check.That(data.ResourceName).Key("data_residency_enabled").HasValue("false"),
 			),
 		},
 		data.ImportStep(),
@@ -76,6 +77,21 @@ func TestAccIotHubDPS_update(t *testing.T) {
 	})
 }
 
+func TestAccIotHubDPS_dataResidencyEnabled(t *testing.T) {
+	data := acceptance.BuildTestData(t, "azurerm_iothub_dps", "test")
+	r := IotHubDPSResource{}
+
+	data.ResourceTest(t, r, []acceptance.TestStep{
+		{
+			Config: r.dataResidencyEnabled(data),
+			Check: acceptance.ComposeTestCheckFunc(
+				check.That(data.ResourceName).ExistsInAzure(r),
+			),
+		},
+		data.ImportStep(),
+	})
+}
+
 func TestAccIotHubDPS_linkedHubs(t *testing.T) {
 	data := acceptance.BuildTestData(t, "azurerm_iothub_dps", "test")
 	r := IotHubDPSResource{}
@@ -228,6 +244,34 @@ resource "azurerm_iothub_dps" "test" {
 `, data.RandomInteger, data.Locations.Primary, data.RandomInteger)
 }
 
+func (IotHubDPSResource) dataResidencyEnabled(data acceptance.TestData) string {
+	// Data residency has limited region support
+	data.Locations.Primary = "brazilsouth"
+	return fmt.Sprintf(`
+provider "azurerm" {
+  features {}
+}
+
+resource "azurerm_resource_group" "test" {
+  name     = "acctestRG-%d"
+  location = "%s"
+}
+
+resource "azurerm_iothub_dps" "test" {
+  name                = "acctestIoTDPS-%d"
+  resource_group_name = azurerm_resource_group.test.name
+  location            = azurerm_resource_group.test.location
+
+  data_residency_enabled = true
+
+  sku {
+    name     = "S1"
+    capacity = "1"
+  }
+}
+`, data.RandomInteger, data.Locations.Primary, data.RandomInteger)
+}
+
 func (IotHubDPSResource) linkedHubs(data acceptance.TestData) string {
 	return fmt.Sprintf(`
 provider "azurerm" {
diff --git a/website/docs/r/iothub_dps.html.markdown b/website/docs/r/iothub_dps.html.markdown
index 757de7295a5c..94d890b52432 100644
--- a/website/docs/r/iothub_dps.html.markdown
+++ b/website/docs/r/iothub_dps.html.markdown
@@ -43,6 +43,8 @@ The following arguments are supported:
 
 * `allocation_policy` - (Optional) The allocation policy of the IoT Device Provisioning Service (`Hashed`, `GeoLatency` or `Static`). Defaults to `Hashed`.
 
+* `data_residency_enabled` - (Optional) Specifies if the IoT Device Provisioning Service has data residency and disaster recovery enabled. Defaults to `false`.
+
 * `sku` - (Required) A `sku` block as defined below.
 
 * `linked_hub` - (Optional) A `linked_hub` block as defined below.