From dab7359650416d5505414259000cc6db358bb757 Mon Sep 17 00:00:00 2001
From: zjhe <hezijie@microsoft.com>
Date: Fri, 9 Dec 2022 08:22:59 +0800
Subject: [PATCH] Make tests work for both `count` and `for_each`

---
 test/e2e/terraform_test.go   | 14 +++++++++++---
 test/upgrade/upgrade_test.go | 14 +++++++++++---
 2 files changed, 22 insertions(+), 6 deletions(-)

diff --git a/test/e2e/terraform_test.go b/test/e2e/terraform_test.go
index 9f9f181..e31aebd 100644
--- a/test/e2e/terraform_test.go
+++ b/test/e2e/terraform_test.go
@@ -1,6 +1,7 @@
 package e2e
 
 import (
+	"fmt"
 	"regexp"
 	"testing"
 
@@ -19,15 +20,22 @@ func TestExamples(t *testing.T) {
 		"examples/private_link_service",
 	}
 	for _, example := range examples {
-		t.Run(example, func(t *testing.T) {
-			testExample(t, example)
+		t.Run(fmt.Sprintf("%s_for_each", example), func(t *testing.T) {
+			testExample(t, example, true)
+		})
+		t.Run(fmt.Sprintf("%s_count", example), func(t *testing.T) {
+			testExample(t, example, false)
 		})
 	}
 }
 
-func testExample(t *testing.T, exampleRelativePath string) {
+func testExample(t *testing.T, exampleRelativePath string, useForEach bool) {
+	vars := map[string]interface{}{
+		"use_for_each": useForEach,
+	}
 	test_helper.RunE2ETest(t, "../../", exampleRelativePath, terraform.Options{
 		Upgrade: true,
+		Vars:    vars,
 	}, func(t *testing.T, output test_helper.TerraformOutput) {
 		vnetId, ok := output["test_vnet_id"].(string)
 		assert.True(t, ok)
diff --git a/test/upgrade/upgrade_test.go b/test/upgrade/upgrade_test.go
index 8bb4c34..9aa75bc 100644
--- a/test/upgrade/upgrade_test.go
+++ b/test/upgrade/upgrade_test.go
@@ -1,6 +1,7 @@
 package upgrade
 
 import (
+	"fmt"
 	"testing"
 
 	test_helper "github.com/Azure/terraform-module-test-helper"
@@ -17,13 +18,16 @@ func TestExamples(t *testing.T) {
 		"examples/private_link_service",
 	}
 	for _, example := range examples {
-		t.Run(example, func(t *testing.T) {
-			testExample(t, example)
+		t.Run(fmt.Sprintf("%s_for_each", example), func(t *testing.T) {
+			testExample(t, example, true)
+		})
+		t.Run(fmt.Sprintf("%s_count", example), func(t *testing.T) {
+			testExample(t, example, false)
 		})
 	}
 }
 
-func testExample(t *testing.T, exampleRelativePath string) {
+func testExample(t *testing.T, exampleRelativePath string, useForEach bool) {
 	currentRoot, err := test_helper.GetCurrentModuleRootPath()
 	if err != nil {
 		t.FailNow()
@@ -32,7 +36,11 @@ func testExample(t *testing.T, exampleRelativePath string) {
 	if err != nil {
 		t.FailNow()
 	}
+	vars := map[string]interface{}{
+		"use_for_each": useForEach,
+	}
 	test_helper.ModuleUpgradeTest(t, "Azure", "terraform-azurerm-vnet", exampleRelativePath, currentRoot, terraform.Options{
 		Upgrade: true,
+		Vars:    vars,
 	}, currentMajorVersion)
 }