Connect nodes with the same state key sequentially

internal/command/test_test.go

@@ -305,9 +305,6 @@ func TestTest_Runs(t *testing.T) {
 		},
 	}
 	for name, tc := range tcs {
-		if name != "shared_state_parallel" {
-			continue
-		}
 		t.Run(name, func(t *testing.T) {
 			if tc.skip {
 				t.Skip()
@@ -506,24 +503,31 @@ func TestTest_Parallel(t *testing.T) {
 	// Split the log into lines
 	lines := strings.Split(output, "\n")
 
-	// Find the positions of "test_d" and "test_c"
-	var testDIndex, testCIndex int
+	// Find the positions of "test_d", "test_c", "test_setup" in the log output
+	var testDIndex, testCIndex, testSetupIndex int
 	for i, line := range lines {
-		if strings.Contains(line, "run \"test_d\"") {
+		if strings.Contains(line, "run \"setup\"") {
+			testSetupIndex = i
+		} else if strings.Contains(line, "run \"test_d\"") {
 			testDIndex = i
 		} else if strings.Contains(line, "run \"test_c\"") {
 			testCIndex = i
 		}
 	}
+	if testDIndex == 0 || testCIndex == 0 || testSetupIndex == 0 {
+		t.Fatalf("test_d, test_c, or test_setup not found in the log output")
+	}
 
 	// Ensure "test_d" appears before "test_c", because test_d has no dependencies,
 	// and would therefore run in parallel to much earlier tests which test_c depends on.
-	if testDIndex == 0 || testCIndex == 0 {
-		t.Fatalf("Could not find both test_d and test_c in the output")
-	}
 	if testDIndex > testCIndex {
 		t.Errorf("test_d appears after test_c in the log output")
 	}
+
+	// Ensure "test_d" appears after "test_setup", because they have the same state key
+	if testDIndex < testSetupIndex {
+		t.Errorf("test_d appears before test_setup in the log output")
+	}
 }
 
 func TestTest_InterruptSkipsRemaining(t *testing.T) {

internal/command/testdata/test/shared_state_parallel/parallel.tftest.hcl

@@ -9,6 +9,7 @@ variables {
 
 run "setup" {
   parallel = true
+  state_key = "test_d"
   module {
     source = "./setup"
   }
@@ -81,6 +82,7 @@ run "test_c" {
 // Does not depend on previous run, and has different state key, so would run in parallel
 // NotDepends: true
 // DiffStateKey: true
+// However, it has a state key that is the same as a previous run, so it should wait for that run.
 run "test_d" {
   parallel = true
   state_key = "test_d"
@@ -93,3 +95,19 @@ run "test_d" {
     error_message = "double bad"
   }
 }
+
+// Does not depend on previous run, and has different state key, so would run in parallel
+// NotDepends: true
+// DiffStateKey: true
+# run "test_d" {
+#   parallel = true
+#   state_key = "test_d"
+#   variables {
+#     input = "foo"
+#   }
+
+#   assert {
+#     condition = output.value == var.foo
+#     error_message = "double bad"
+#   }
+# }

internal/terraformtest/node_test_run.go

@@ -6,14 +6,12 @@ package terraformtest
 import (
 	"fmt"
 
-	"github.com/hashicorp/terraform/internal/configs"
 	"github.com/hashicorp/terraform/internal/moduletest"
 )
 
 type NodeTestRun struct {
-	file   *moduletest.File
-	run    *moduletest.Run
-	config *configs.Config
+	file *moduletest.File
+	run  *moduletest.Run
 }
 
 func (n *NodeTestRun) Run() *moduletest.Run {

internal/terraformtest/transform_test_run.go

@@ -26,27 +26,45 @@ type TestRunTransformer struct {
 }
 
 func (t *TestRunTransformer) Transform(g *terraform.Graph) error {
-	var prev *NodeTestRun
 	var errs []error
-	runsSoFar := make(map[string]*NodeTestRun)
-	for _, run := range t.File.Runs {
-		// If we're testing a specific configuration, we need to use that
-		config := t.config
-		if run.Config.ConfigUnderTest != nil {
-			config = run.Config.ConfigUnderTest
-		}
 
-		node := &NodeTestRun{run: run, file: t.File, config: config}
-		g.Add(node)
+	// Create and add nodes for each run
+	nodes, err := t.createNodes(g)
+	if err != nil {
+		return err
+	}
 
-		// parallelized sequential runs are only connected if they have the same state key or if they depend on each other
-		if prev != nil && prev.run.GetStateKey() == run.GetStateKey() && prev.run.Config.Parallel && run.Config.Parallel {
-			g.Connect(dag.BasicEdge(node, prev))
-		}
-		prev = node
+	// Connect nodes based on dependencies
+	if err := t.connectDependencies(g, nodes); err != nil {
+		errs = append(errs, err)
+	}
 
-		// Connect the run to all the other runs that it depends on
-		refs, err := getRefs(run)
+	// Connect nodes with the same state key sequentially
+	if err := t.connectStateKeyRuns(g, nodes); err != nil {
+		errs = append(errs, err)
+	}
+
+	return errors.Join(errs...)
+}
+
+func (t *TestRunTransformer) createNodes(g *terraform.Graph) ([]*NodeTestRun, error) {
+	var nodes []*NodeTestRun
+	for _, run := range t.File.Runs {
+		node := &NodeTestRun{run: run, file: t.File}
+		g.Add(node)
+		nodes = append(nodes, node)
+	}
+	return nodes, nil
+}
+
+func (t *TestRunTransformer) connectDependencies(g *terraform.Graph, nodes []*NodeTestRun) error {
+	var errs []error
+	nodeMap := make(map[string]*NodeTestRun)
+	for _, node := range nodes {
+		nodeMap[node.run.Name] = node
+	}
+	for _, node := range nodes {
+		refs, err := getRefs(node.run)
 		if err != nil {
 			return err
 		}
@@ -59,19 +77,31 @@ func (t *TestRunTransformer) Transform(g *terraform.Graph) error {
 			if runName == "" {
 				continue
 			}
-			dependency, ok := runsSoFar[runName]
+			dependency, ok := nodeMap[runName]
 			if !ok {
-				errs = append(errs, fmt.Errorf("dependency `run.%s` not found for run %q", runName, run.Name))
+				errs = append(errs, fmt.Errorf("dependency `run.%s` not found for run %q", runName, node.run.Name))
 				continue
 			}
 			g.Connect(dag.BasicEdge(node, dependency))
 		}
-		runsSoFar[run.Name] = node
 	}
-
 	return errors.Join(errs...)
 }
 
+func (t *TestRunTransformer) connectStateKeyRuns(g *terraform.Graph, nodes []*NodeTestRun) error {
+	stateRuns := make(map[string][]*NodeTestRun)
+	for _, node := range nodes {
+		key := node.run.GetStateKey()
+		stateRuns[key] = append(stateRuns[key], node)
+	}
+	for _, runs := range stateRuns {
+		for i := 1; i < len(runs); i++ {
+			g.Connect(dag.BasicEdge(runs[i], runs[i-1]))
+		}
+	}
+	return nil
+}
+
 func getRefs(run *moduletest.Run) ([]*addrs.Reference, error) {
 	refs, refDiags := run.GetReferences()
 	if refDiags.HasErrors() {