Support health extension for rolling ugrade mode

azurerm/internal/services/compute/linux_virtual_machine_scale_set_resource.go

@@ -332,10 +332,6 @@ func resourceArmLinuxVirtualMachineScaleSetCreate(d *schema.ResourceData, meta i
 		return fmt.Errorf("An `automatic_os_upgrade_policy` block cannot be specified when `upgrade_mode` is not set to `Automatic`")
 	}
 
-	if upgradeMode == compute.UpgradeModeAutomatic && len(automaticOSUpgradePolicyRaw) > 0 && healthProbeId == "" {
-		return fmt.Errorf("`healthProbeId` must be set when `upgrade_mode` is set to %q and `automatic_os_upgrade_policy` block exists", string(upgradeMode))
-	}
-
 	shouldHaveRollingUpgradePolicy := upgradeMode == compute.UpgradeModeAutomatic || upgradeMode == compute.UpgradeModeRolling
 	if !shouldHaveRollingUpgradePolicy && len(rollingUpgradePolicyRaw) > 0 {
 		return fmt.Errorf("A `rolling_upgrade_policy` block cannot be specified when `upgrade_mode` is set to %q", string(upgradeMode))
@@ -402,15 +398,20 @@ func resourceArmLinuxVirtualMachineScaleSetCreate(d *schema.ResourceData, meta i
 		},
 	}
 
+	hasHealthExtension := false
 	if features.VMSSExtensionsBeta() {
 		if vmExtensionsRaw, ok := d.GetOk("extension"); ok {
-			virtualMachineProfile.ExtensionProfile, err = expandVirtualMachineScaleSetExtensions(vmExtensionsRaw.([]interface{}))
+			virtualMachineProfile.ExtensionProfile, hasHealthExtension, err = expandVirtualMachineScaleSetExtensions(vmExtensionsRaw.([]interface{}))
 			if err != nil {
 				return err
 			}
 		}
 	}
 
+	if upgradeMode == compute.UpgradeModeAutomatic && len(automaticOSUpgradePolicyRaw) > 0 && (healthProbeId == "" || !hasHealthExtension) {
+		return fmt.Errorf("`health_probe_id` must be set or a health extension must be specified when `upgrade_mode` is set to %q and `automatic_os_upgrade_policy` block exists", string(upgradeMode))
+	}
+
 	if adminPassword, ok := d.GetOk("admin_password"); ok {
 		virtualMachineProfile.OsProfile.AdminPassword = utils.String(adminPassword.(string))
 	}
@@ -774,7 +775,7 @@ func resourceArmLinuxVirtualMachineScaleSetUpdate(d *schema.ResourceData, meta i
 		if d.HasChange("extension") {
 			updateInstances = true
 
-			extensionProfile, err := expandVirtualMachineScaleSetExtensions(d.Get("extension").([]interface{}))
+			extensionProfile, _, err := expandVirtualMachineScaleSetExtensions(d.Get("extension").([]interface{}))
 			if err != nil {
 				return err
 			}

azurerm/internal/services/compute/tests/linux_virtual_machine_scale_set_other_resource_test.go

@@ -705,6 +705,25 @@ func TestAccAzureRMLinuxVirtualMachineScaleSet_otherVmExtensionsUpdate(t *testin
 	})
 }
 
+func TestAccAzureRMLinuxVirtualMachineScaleSet_otherVmExtensionsRollingUpgradeWithHealthExtension(t *testing.T) {
+	data := acceptance.BuildTestData(t, "azurerm_linux_virtual_machine_scale_set", "test")
+
+	resource.ParallelTest(t, resource.TestCase{
+		PreCheck:     func() { acceptance.PreCheck(t) },
+		Providers:    acceptance.SupportedProviders,
+		CheckDestroy: testCheckAzureRMLinuxVirtualMachineScaleSetDestroy,
+		Steps: []resource.TestStep{
+			{
+				Config: testAccAzureRMLinuxVirtualMachineScaleSet_otherVmExtensionsRollingUpgradeWithHealthExtension(data),
+				Check: resource.ComposeTestCheckFunc(
+					testCheckAzureRMLinuxVirtualMachineScaleSetExists(data.ResourceName)),
+			},
+			// TODO - extension should be changed to extension.0.protected_settings when either binary testing is available or this feature is promoted from beta
+			data.ImportStep("admin_password", "extension"),
+		},
+	})
+}
+
 func TestAccAzureRMLinuxVirtualMachineScaleSet_otherEncryptionAtHost(t *testing.T) {
 	data := acceptance.BuildTestData(t, "azurerm_linux_virtual_machine_scale_set", "test")
 
@@ -2179,7 +2198,6 @@ resource "azurerm_linux_virtual_machine_scale_set" "test" {
     settings = jsonencode({
       "commandToExecute" = "echo $HOSTNAME"
     })
-
   }
 
   tags = {
@@ -2246,7 +2264,6 @@ resource "azurerm_linux_virtual_machine_scale_set" "test" {
     protected_settings = jsonencode({
       "managedIdentity" = {}
     })
-
   }
 
   tags = {
@@ -2382,7 +2399,6 @@ resource "azurerm_linux_virtual_machine_scale_set" "test" {
     protected_settings = jsonencode({
       "managedIdentity" = {}
     })
-
   }
 
   extension {
@@ -2405,6 +2421,79 @@ resource "azurerm_linux_virtual_machine_scale_set" "test" {
 `, template, data.RandomInteger)
 }
 
+func testAccAzureRMLinuxVirtualMachineScaleSet_otherVmExtensionsRollingUpgradeWithHealthExtension(data acceptance.TestData) string {
+	template := testAccAzureRMLinuxVirtualMachineScaleSet_template(data)
+	return fmt.Sprintf(`
+%[1]s
+
+provider "azurerm" {
+  features {}
+}
+
+resource "azurerm_linux_virtual_machine_scale_set" "test" {
+  name                = "acctestvmss-%d"
+  resource_group_name = azurerm_resource_group.test.name
+  location            = azurerm_resource_group.test.location
+  sku                 = "Standard_F2"
+  instances           = 1
+  admin_username      = "adminuser"
+  admin_password      = "P@ssword1234!"
+
+  disable_password_authentication = false
+
+  upgrade_mode = "Rolling"
+
+  rolling_upgrade_policy {
+    max_batch_instance_percent              = 21
+    max_unhealthy_instance_percent          = 22
+    max_unhealthy_upgraded_instance_percent = 23
+    pause_time_between_batches              = "PT30S"
+  }
+
+  source_image_reference {
+    publisher = "Canonical"
+    offer     = "UbuntuServer"
+    sku       = "16.04-LTS"
+    version   = "latest"
+  }
+
+  os_disk {
+    storage_account_type = "Standard_LRS"
+    caching              = "ReadWrite"
+  }
+
+  network_interface {
+    name    = "example"
+    primary = true
+
+    ip_configuration {
+      name      = "internal"
+      primary   = true
+      subnet_id = azurerm_subnet.test.id
+    }
+  }
+
+  extension {
+    name                       = "HealthExtension"
+    publisher                  = "Microsoft.ManagedServices"
+    type                       = "ApplicationHealthLinux"
+    type_handler_version       = "1.0"
+    auto_upgrade_minor_version = true
+
+    settings = jsonencode({
+      protocol     = "https"
+      port         = 443
+      request_path = "/"
+    })
+  }
+
+  tags = {
+    accTest = "true"
+  }
+}
+`, template, data.RandomInteger)
+}
+
 func testAccAzureRMLinuxVirtualMachineScaleSet_otherVmExtensionsUpdate(data acceptance.TestData) string {
 	template := testAccAzureRMLinuxVirtualMachineScaleSet_template(data)
 	return fmt.Sprintf(`

azurerm/internal/services/compute/tests/windows_virtual_machine_scale_set_other_resource_test.go

@@ -864,6 +864,26 @@ func TestAccAzureRMWindowsVirtualMachineScaleSet_otherVmExtensionUpdate(t *testi
 	})
 }
 
+func TestAccAzureRMWindowsVirtualMachineScaleSet_otherVmExtensionsRollingUpgradeWithHealthExtension(t *testing.T) {
+	data := acceptance.BuildTestData(t, "azurerm_windows_virtual_machine_scale_set", "test")
+
+	resource.ParallelTest(t, resource.TestCase{
+		PreCheck:     func() { acceptance.PreCheck(t) },
+		Providers:    acceptance.SupportedProviders,
+		CheckDestroy: testCheckAzureRMWindowsVirtualMachineScaleSetDestroy,
+		Steps: []resource.TestStep{
+			{
+				Config: testAccAzureRMWindowsVirtualMachineScaleSet_otherVmExtensionsRollingUpgradeWithHealthExtension(data),
+				Check: resource.ComposeTestCheckFunc(
+					testCheckAzureRMWindowsVirtualMachineScaleSetExists(data.ResourceName),
+				),
+			},
+			// TODO - extension should be changed to extension.0.protected_settings when either binary testing is available or this feature is promoted from beta
+			data.ImportStep("admin_password", "extension"),
+		},
+	})
+}
+
 func TestAccAzureRMWindowsVirtualMachineScaleSet_otherEncryptionAtHostEnabled(t *testing.T) {
 	data := acceptance.BuildTestData(t, "azurerm_windows_virtual_machine_scale_set", "test")
 
@@ -2740,6 +2760,73 @@ resource "azurerm_windows_virtual_machine_scale_set" "test" {
 `, template)
 }
 
+func testAccAzureRMWindowsVirtualMachineScaleSet_otherVmExtensionsRollingUpgradeWithHealthExtension(data acceptance.TestData) string {
+	template := testAccAzureRMWindowsVirtualMachineScaleSet_template(data)
+	return fmt.Sprintf(`
+%s
+
+provider "azurerm" {
+  features {}
+}
+
+resource "azurerm_windows_virtual_machine_scale_set" "test" {
+  name                = local.vm_name
+  resource_group_name = azurerm_resource_group.test.name
+  location            = azurerm_resource_group.test.location
+  sku                 = "Standard_F2"
+  instances           = 1
+  admin_username      = "adminuser"
+  admin_password      = "P@ssword1234!"
+
+  upgrade_mode = "Rolling"
+
+  rolling_upgrade_policy {
+    max_batch_instance_percent              = 21
+    max_unhealthy_instance_percent          = 22
+    max_unhealthy_upgraded_instance_percent = 23
+    pause_time_between_batches              = "PT30S"
+  }
+
+  source_image_reference {
+    publisher = "MicrosoftWindowsServer"
+    offer     = "WindowsServer"
+    sku       = "2019-Datacenter"
+    version   = "latest"
+  }
+
+  os_disk {
+    storage_account_type = "Standard_LRS"
+    caching              = "ReadWrite"
+  }
+
+  network_interface {
+    name    = "example"
+    primary = true
+
+    ip_configuration {
+      name      = "internal"
+      primary   = true
+      subnet_id = azurerm_subnet.test.id
+    }
+  }
+
+  extension {
+    name                       = "HealthExtension"
+    publisher                  = "Microsoft.ManagedServices"
+    type                       = "ApplicationHealthWindows"
+    type_handler_version       = "1.0"
+    auto_upgrade_minor_version = true
+
+    settings = jsonencode({
+      protocol     = "https"
+      port         = 443
+      request_path = "/"
+    })
+  }
+}
+`, template)
+}
+
 func testAccAzureRMWindowsVirtualMachineScaleSet_otherUpgradeMode(data acceptance.TestData, enabled bool) string {
 	template := testAccAzureRMWindowsVirtualMachineScaleSet_template(data)
 	return fmt.Sprintf(`

azurerm/internal/services/compute/virtual_machine_scale_set.go

@@ -1367,27 +1367,33 @@ func VirtualMachineScaleSetExtensionsSchema() *schema.Schema {
 	}
 }
 
-func expandVirtualMachineScaleSetExtensions(input []interface{}) (*compute.VirtualMachineScaleSetExtensionProfile, error) {
+func expandVirtualMachineScaleSetExtensions(input []interface{}) (*compute.VirtualMachineScaleSetExtensionProfile, bool, error) {
 	result := &compute.VirtualMachineScaleSetExtensionProfile{}
 	if len(input) == 0 {
-		return result, nil
+		return result, false, nil
 	}
 
 	extensions := make([]compute.VirtualMachineScaleSetExtension, 0)
+	hasHealthExtension := false
 	for _, v := range input {
 		extensionRaw := v.(map[string]interface{})
 		extension := compute.VirtualMachineScaleSetExtension{
 			Name: utils.String(extensionRaw["name"].(string)),
 		}
+		extensionType := extensionRaw["type"].(string)
 
 		extensionProps := compute.VirtualMachineScaleSetExtensionProperties{
 			Publisher:                utils.String(extensionRaw["publisher"].(string)),
-			Type:                     utils.String(extensionRaw["type"].(string)),
+			Type:                     &extensionType,
 			TypeHandlerVersion:       utils.String(extensionRaw["type_handler_version"].(string)),
 			AutoUpgradeMinorVersion:  utils.Bool(extensionRaw["auto_upgrade_minor_version"].(bool)),
 			ProvisionAfterExtensions: utils.ExpandStringSlice(extensionRaw["provision_after_extensions"].([]interface{})),
 		}
 
+		if extensionType == "ApplicationHealthLinux" || extensionType == "ApplicationHealthWindows" {
+			hasHealthExtension = true
+		}
+
 		if forceUpdateTag := extensionRaw["force_update_tag"]; forceUpdateTag != nil {
 			extensionProps.ForceUpdateTag = utils.String(forceUpdateTag.(string))
 		}
@@ -1396,7 +1402,7 @@ func expandVirtualMachineScaleSetExtensions(input []interface{}) (*compute.Virtu
 		if ok && settings != "" {
 			settings, err := structure.ExpandJsonFromString(settings)
 			if err != nil {
-				return nil, fmt.Errorf("failed to parse JSON from `settings`: %+v", err)
+				return nil, false, fmt.Errorf("failed to parse JSON from `settings`: %+v", err)
 			}
 			extensionProps.Settings = settings
 		}
@@ -1405,7 +1411,7 @@ func expandVirtualMachineScaleSetExtensions(input []interface{}) (*compute.Virtu
 		if ok && protectedSettings != "" {
 			protectedSettings, err := structure.ExpandJsonFromString(protectedSettings)
 			if err != nil {
-				return nil, fmt.Errorf("failed to parse JSON from `settings`: %+v", err)
+				return nil, false, fmt.Errorf("failed to parse JSON from `settings`: %+v", err)
 			}
 			extensionProps.ProtectedSettings = protectedSettings
 		}
@@ -1415,7 +1421,7 @@ func expandVirtualMachineScaleSetExtensions(input []interface{}) (*compute.Virtu
 	}
 	result.Extensions = &extensions
 
-	return result, nil
+	return result, hasHealthExtension, nil
 }
 
 func flattenVirtualMachineScaleSetExtensions(input *compute.VirtualMachineScaleSetExtensionProfile, d *schema.ResourceData) ([]map[string]interface{}, error) {

azurerm/internal/services/compute/windows_virtual_machine_scale_set_resource.go

@@ -349,10 +349,6 @@ func resourceArmWindowsVirtualMachineScaleSetCreate(d *schema.ResourceData, meta
 		return fmt.Errorf("An `automatic_os_upgrade_policy` block cannot be specified when `upgrade_mode` is not set to `Automatic`")
 	}
 
-	if upgradeMode == compute.UpgradeModeAutomatic && len(automaticOSUpgradePolicyRaw) > 0 && healthProbeId == "" {
-		return fmt.Errorf("`healthProbeId` must be set when `upgrade_mode` is set to %q and `automatic_os_upgrade_policy` block exists", string(upgradeMode))
-	}
-
 	shouldHaveRollingUpgradePolicy := upgradeMode == compute.UpgradeModeAutomatic || upgradeMode == compute.UpgradeModeRolling
 	if !shouldHaveRollingUpgradePolicy && len(rollingUpgradePolicyRaw) > 0 {
 		return fmt.Errorf("A `rolling_upgrade_policy` block cannot be specified when `upgrade_mode` is set to %q", string(upgradeMode))
@@ -419,15 +415,20 @@ func resourceArmWindowsVirtualMachineScaleSetCreate(d *schema.ResourceData, meta
 		},
 	}
 
+	hasHealthExtension := false
 	if features.VMSSExtensionsBeta() {
 		if vmExtensionsRaw, ok := d.GetOk("extension"); ok {
-			virtualMachineProfile.ExtensionProfile, err = expandVirtualMachineScaleSetExtensions(vmExtensionsRaw.([]interface{}))
+			virtualMachineProfile.ExtensionProfile, hasHealthExtension, err = expandVirtualMachineScaleSetExtensions(vmExtensionsRaw.([]interface{}))
 			if err != nil {
 				return err
 			}
 		}
 	}
 
+	if upgradeMode == compute.UpgradeModeAutomatic && len(automaticOSUpgradePolicyRaw) > 0 && (healthProbeId == "" || !hasHealthExtension) {
+		return fmt.Errorf("`health_probe_id` must be set or a health extension must be specified when `upgrade_mode` is set to %q and `automatic_os_upgrade_policy` block exists", string(upgradeMode))
+	}
+
 	enableAutomaticUpdates := d.Get("enable_automatic_updates").(bool)
 	if upgradeMode != compute.UpgradeModeAutomatic {
 		virtualMachineProfile.OsProfile.WindowsConfiguration.EnableAutomaticUpdates = utils.Bool(enableAutomaticUpdates)
@@ -804,7 +805,7 @@ func resourceArmWindowsVirtualMachineScaleSetUpdate(d *schema.ResourceData, meta
 
 	if features.VMSSExtensionsBeta() {
 		if d.HasChange("extension") {
-			extensionProfile, err := expandVirtualMachineScaleSetExtensions(d.Get("extension").([]interface{}))
+			extensionProfile, _, err := expandVirtualMachineScaleSetExtensions(d.Get("extension").([]interface{}))
 			if err != nil {
 				return err
 			}