stackstate: A helper for loading state during tests

The main entry point here assumes the caller is providing serialized state
objects wrapped in anypb.Any messages.

That's an inconvenient representation for hand-writing in tests, so the
new function LoadFromDirectProto allows skipping the deserialization steps
and instead has the caller write the wanted values as if they had already
been parsed/unmarshaled.

internal/stacks/stackruntime/internal/stackeval/testing_test.go

@@ -11,8 +11,10 @@ import (
 	"github.com/hashicorp/terraform/internal/stacks/stackaddrs"
 	"github.com/hashicorp/terraform/internal/stacks/stackconfig"
 	"github.com/hashicorp/terraform/internal/stacks/stackstate"
+	"github.com/hashicorp/terraform/internal/stacks/stackstate/statekeys"
 	"github.com/hashicorp/terraform/internal/tfdiags"
 	"github.com/zclconf/go-cty/cty"
+	"google.golang.org/protobuf/reflect/protoreflect"
 )
 
 // This file contains some general test utilities that many of our other
@@ -61,6 +63,15 @@ func testSourceBundle(t *testing.T) *sourcebundle.Bundle {
 	return sources
 }
 
+func testPriorState(t *testing.T, msgs map[statekeys.Key]protoreflect.ProtoMessage) *stackstate.State {
+	t.Helper()
+	ret, err := stackstate.LoadFromDirectProto(msgs)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return ret
+}
+
 // testEvaluator constructs a [Main] that's configured for [InspectPhase] using
 // the given configuration, state, and other options.
 //

internal/stacks/stackstate/from_proto.go

@@ -16,6 +16,10 @@ import (
 	"google.golang.org/protobuf/types/known/anypb"
 )
 
+// LoadFromProto produces a [State] object by decoding a raw state map.
+//
+// This is the primary way to load a "prior state" provided by a caller
+// into memory so we can use it in the stack runtime.
 func LoadFromProto(msgs map[string]*anypb.Any) (*State, error) {
 	ret := NewState()
 	ret.inputRaw = msgs
@@ -30,47 +34,9 @@ func LoadFromProto(msgs map[string]*anypb.Any) (*State, error) {
 		}
 
 		if !statekeys.RecognizedType(key) {
-			// There are three different strategies for dealing with
-			// unrecognized keys, which we recognize based on naming
-			// conventions of the key types.
-			switch handling := key.KeyType().UnrecognizedKeyHandling(); handling {
-
-			case statekeys.FailIfUnrecognized:
-				// This is for keys whose messages materially change the
-				// meaning of the state and so cannot be ignored. Keys
-				// with this treatment are forwards-incompatible (old versions
-				// of Terraform will fail to load a state containing them) so
-				// should be added only as a last resort.
-				return nil, fmt.Errorf("state was created by a newer version of Terraform Core (unrecognized tracking key %q)", rawKey)
-
-			case statekeys.PreserveIfUnrecognized:
-				// This is for keys whose messages can safely be left entirely
-				// unchanged if applying a plan with a version of Terraform
-				// that doesn't understand them. Keys in this category should
-				// typically be standalone and not refer to or depend on any
-				// other objects in the state, to ensure that removing or
-				// updating other objects will not cause the preserved message
-				// to become misleading or invalid.
-				// We don't need to do anything special with these ones because
-				// the caller should preserve any object we don't explicitly
-				// update or delete during the apply phase.
-
-			case statekeys.DiscardIfUnrecognized:
-				// This is for keys which can be discarded when planning or
-				// applying with an older version of Terraform that doesn't
-				// understand them. This category is for optional ancillary
-				// information -- not actually required for correct subsequent
-				// planning -- especially if it could be recomputed again and
-				// repopulated if later planning and applying with a newer
-				// version of Terraform Core.
-				// For these ones we need to remember their keys so that we
-				// can emit "delete" messages early in the apply phase to
-				// actually discard them from the caller's records.
-				ret.discardUnsupportedKeys.Add(key)
-
-			default:
-				// Should not get here. The above should be exhaustive.
-				panic(fmt.Sprintf("unsupported UnrecognizedKeyHandling value %s", handling))
+			err = handleUnrecognizedKey(key, ret)
+			if err != nil {
+				return nil, err
 			}
 			continue
 		}
@@ -80,29 +46,116 @@ func LoadFromProto(msgs map[string]*anypb.Any) (*State, error) {
 			return nil, fmt.Errorf("invalid raw value for raw state key %q: %w", rawKey, err)
 		}
 
-		switch key := key.(type) {
-
-		case statekeys.ComponentInstance:
-			err := handleComponentInstanceMsg(key, msg, ret)
-			if err != nil {
-				return nil, err
-			}
+		err = handleProtoMsg(key, msg, ret)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return ret, nil
+}
 
-		case statekeys.ResourceInstanceObject:
-			err := handleResourceInstanceObjectMsg(key, msg, ret)
+// LoadFromDirectProto is a variation of the primary entry-point [LoadFromProto]
+// which accepts direct messages of the relevant types from the tfstackdata1
+// package, rather than the [anypb.Raw] representation thereof.
+//
+// This is primarily for internal testing purposes, where it's typically more
+// convenient to write out a struct literal for one of the message types
+// directly rather than having to first serialize it to [anypb.Any] only for
+// it to be unserialized again promptly afterwards.
+//
+// Unlike [LoadFromProto], the state object produced by this function will not
+// have any record of the "raw state" it was created from, because this function
+// is bypassing the concept of raw state. [State.InputRaw] will therefore
+// return an empty map.
+//
+// Prefer to use [LoadFromProto] when processing user input. This function
+// cannot accept [anypb.Any] messages even though the Go compiler can't enforce
+// that at compile time.
+func LoadFromDirectProto(msgs map[statekeys.Key]protoreflect.ProtoMessage) (*State, error) {
+	ret := NewState()
+	ret.inputRaw = nil // this doesn't get populated by this entry point
+	for key, msg := range msgs {
+		// The following should be equivalent to the similar loop in
+		// [LoadFromProto] except for skipping the parsing/unmarshalling
+		// steps since key and msg are already in their in-memory forms.
+		if !statekeys.RecognizedType(key) {
+			err := handleUnrecognizedKey(key, ret)
 			if err != nil {
 				return nil, err
 			}
-
-		default:
-			// Should not get here: the above should be exhaustive for all
-			// possible key types.
-			panic(fmt.Sprintf("unsupported state key type %T", key))
+			continue
+		}
+		err := handleProtoMsg(key, msg, ret)
+		if err != nil {
+			return nil, err
 		}
 	}
 	return ret, nil
 }
 
+func handleUnrecognizedKey(key statekeys.Key, state *State) error {
+	// There are three different strategies for dealing with
+	// unrecognized keys, which we recognize based on naming
+	// conventions of the key types.
+	switch handling := key.KeyType().UnrecognizedKeyHandling(); handling {
+
+	case statekeys.FailIfUnrecognized:
+		// This is for keys whose messages materially change the
+		// meaning of the state and so cannot be ignored. Keys
+		// with this treatment are forwards-incompatible (old versions
+		// of Terraform will fail to load a state containing them) so
+		// should be added only as a last resort.
+		return fmt.Errorf("state was created by a newer version of Terraform Core (unrecognized tracking key %q)", statekeys.String(key))
+
+	case statekeys.PreserveIfUnrecognized:
+		// This is for keys whose messages can safely be left entirely
+		// unchanged if applying a plan with a version of Terraform
+		// that doesn't understand them. Keys in this category should
+		// typically be standalone and not refer to or depend on any
+		// other objects in the state, to ensure that removing or
+		// updating other objects will not cause the preserved message
+		// to become misleading or invalid.
+		// We don't need to do anything special with these ones because
+		// the caller should preserve any object we don't explicitly
+		// update or delete during the apply phase.
+		return nil
+
+	case statekeys.DiscardIfUnrecognized:
+		// This is for keys which can be discarded when planning or
+		// applying with an older version of Terraform that doesn't
+		// understand them. This category is for optional ancillary
+		// information -- not actually required for correct subsequent
+		// planning -- especially if it could be recomputed again and
+		// repopulated if later planning and applying with a newer
+		// version of Terraform Core.
+		// For these ones we need to remember their keys so that we
+		// can emit "delete" messages early in the apply phase to
+		// actually discard them from the caller's records.
+		state.discardUnsupportedKeys.Add(key)
+		return nil
+
+	default:
+		// Should not get here. The above should be exhaustive.
+		panic(fmt.Sprintf("unsupported UnrecognizedKeyHandling value %s", handling))
+	}
+}
+
+func handleProtoMsg(key statekeys.Key, msg protoreflect.ProtoMessage, state *State) error {
+	switch key := key.(type) {
+
+	case statekeys.ComponentInstance:
+		return handleComponentInstanceMsg(key, msg, state)
+
+	case statekeys.ResourceInstanceObject:
+		return handleResourceInstanceObjectMsg(key, msg, state)
+
+	default:
+		// Should not get here: the above should be exhaustive for all
+		// possible key types.
+		panic(fmt.Sprintf("unsupported state key type %T", key))
+	}
+}
+
 func handleComponentInstanceMsg(key statekeys.ComponentInstance, msg protoreflect.ProtoMessage, state *State) error {
 	// For this particular object type all of the information is in the key,
 	// for now at least.