From b15a068e2760326612c5d0d3a5a72bae1b08b90c Mon Sep 17 00:00:00 2001
From: Fred Park <fred.park@microsoft.com>
Date: Fri, 5 May 2023 18:06:38 +0000
Subject: [PATCH] azurerm_batch_account - allow versionless keys for CMK

---
 .../services/batch/batch_account_resource.go  |  2 +-
 .../batch/batch_account_resource_test.go      | 72 ++++++++++++++++++-
 website/docs/r/batch_account.html.markdown    |  2 +-
 3 files changed, 71 insertions(+), 5 deletions(-)

diff --git a/internal/services/batch/batch_account_resource.go b/internal/services/batch/batch_account_resource.go
index a7cf1fa6589b..6d7f53887bb7 100644
--- a/internal/services/batch/batch_account_resource.go
+++ b/internal/services/batch/batch_account_resource.go
@@ -159,7 +159,7 @@ func resourceBatchAccount() *pluginsdk.Resource {
 						"key_vault_key_id": {
 							Type:         pluginsdk.TypeString,
 							Required:     true,
-							ValidateFunc: keyVaultValidate.NestedItemId,
+							ValidateFunc: keyVaultValidate.NestedItemIdWithOptionalVersion,
 						},
 					},
 				},
diff --git a/internal/services/batch/batch_account_resource_test.go b/internal/services/batch/batch_account_resource_test.go
index c5f51f1e2482..15140fcaef46 100644
--- a/internal/services/batch/batch_account_resource_test.go
+++ b/internal/services/batch/batch_account_resource_test.go
@@ -158,14 +158,30 @@ func TestAccBatchAccount_identity(t *testing.T) {
 	})
 }
 
-func TestAccBatchAccount_cmk(t *testing.T) {
+func TestAccBatchAccount_cmkVersionedKey(t *testing.T) {
 	data := acceptance.BuildTestData(t, "azurerm_batch_account", "test")
 	r := BatchAccountResource{}
 	tenantID := os.Getenv("ARM_TENANT_ID")
 
 	data.ResourceTest(t, r, []acceptance.TestStep{
 		{
-			Config: r.cmk(data, tenantID),
+			Config: r.cmkVersionedKey(data, tenantID),
+			Check: acceptance.ComposeTestCheckFunc(
+				check.That(data.ResourceName).ExistsInAzure(r),
+				check.That(data.ResourceName).Key("encryption.0.key_vault_key_id").IsSet(),
+			),
+		},
+	})
+}
+
+func TestAccBatchAccount_cmkVersionlessKey(t *testing.T) {
+	data := acceptance.BuildTestData(t, "azurerm_batch_account", "test")
+	r := BatchAccountResource{}
+	tenantID := os.Getenv("ARM_TENANT_ID")
+
+	data.ResourceTest(t, r, []acceptance.TestStep{
+		{
+			Config: r.cmkVersionlessKey(data, tenantID),
 			Check: acceptance.ComposeTestCheckFunc(
 				check.That(data.ResourceName).ExistsInAzure(r),
 				check.That(data.ResourceName).Key("encryption.0.key_vault_key_id").IsSet(),
@@ -543,7 +559,7 @@ resource "azurerm_batch_account" "test" {
 `, data.RandomInteger, data.Locations.Primary, data.RandomString, data.RandomString, data.RandomString)
 }
 
-func (BatchAccountResource) cmk(data acceptance.TestData, tenantID string) string {
+func (BatchAccountResource) cmkVersionedKey(data acceptance.TestData, tenantID string) string {
 	return fmt.Sprintf(`
 provider "azurerm" {
   features {
@@ -632,6 +648,56 @@ resource "azurerm_key_vault" "test" {
   }
 }
 
+func (BatchAccountResource) cmkVersionlessKey(data acceptance.TestData, tenantID string) string {
+	return fmt.Sprintf(`
+provider "azurerm" {
+  features {
+    key_vault {
+      purge_soft_delete_on_destroy       = false
+      purge_soft_deleted_keys_on_destroy = false
+    }
+  }
+}
+
+data "azurerm_client_config" "current" {
+}
+
+resource "azurerm_resource_group" "test" {
+  name     = "testaccRG-batch-%d"
+  location = "%s"
+}
+
+resource "azurerm_storage_account" "test" {
+  name                     = "testaccsa%s"
+  resource_group_name      = azurerm_resource_group.test.name
+  location                 = azurerm_resource_group.test.location
+  account_tier             = "Standard"
+  account_replication_type = "LRS"
+}
+
+resource "azurerm_user_assigned_identity" "test" {
+  name                = "acctest%s"
+  resource_group_name = azurerm_resource_group.test.name
+  location            = azurerm_resource_group.test.location
+}
+
+resource "azurerm_batch_account" "test" {
+  name                                = "testaccbatch%s"
+  resource_group_name                 = azurerm_resource_group.test.name
+  location                            = azurerm_resource_group.test.location
+  pool_allocation_mode                = "BatchService"
+  storage_account_id                  = azurerm_storage_account.test.id
+  storage_account_authentication_mode = "StorageKeys"
+  identity {
+    type         = "UserAssigned"
+    identity_ids = [azurerm_user_assigned_identity.test.id]
+  }
+
+  encryption {
+    key_vault_key_id = "${azurerm_key_vault.test.vault_uri}keys/${azurerm_key_vault_key.test.name}"
+  }
+}
+
 resource "azurerm_key_vault_key" "test" {
   name         = "enckey%d"
   key_vault_id = "${azurerm_key_vault.test.id}"
diff --git a/website/docs/r/batch_account.html.markdown b/website/docs/r/batch_account.html.markdown
index f0194926114f..aec71915805c 100644
--- a/website/docs/r/batch_account.html.markdown
+++ b/website/docs/r/batch_account.html.markdown
@@ -103,7 +103,7 @@ A `key_vault_reference` block supports the following:
 
 A `encryption` block supports the following:
 
-* `key_vault_key_id` - (Required) The Azure key vault reference id with version that should be used to encrypt data, as documented [here](https://docs.microsoft.com/azure/batch/batch-customer-managed-key). Key rotation is not yet supported.
+* `key_vault_key_id` - (Required) The Azure key vault reference id that should be used to encrypt data, as documented [here](https://docs.microsoft.com/azure/batch/batch-customer-managed-key). Both versioned and versionless keys are supported.
 
 ## Attributes Reference