Consolidate provider resource mocking and overrides (#3547)

Signed-off-by: Christian Mesh <christianmesh1@gmail.com>

internal/command/e2etest/provider_plugin_test.go

@@ -129,12 +129,12 @@ func TestProviderGlobalCache(t *testing.T) {
 	for i := 0; i < 16; i++ {
 		wg.Add(1)
 		go func() {
+			defer wg.Done()
 			tf := e2e.NewBinary(t, tofuBin, "testdata/provider-global-cache")
 			tf.AddEnv(fmt.Sprintf("TF_CLI_CONFIG_FILE=%s", rcLoc))
 
 			stdout, stderr, err := tf.Run("init")
 			tofuResult{t, stdout, stderr, err}.Success()
-			wg.Done()
 		}()
 	}
 

internal/tofu/eval_context_builtin.go

@@ -148,22 +148,6 @@ func (c *BuiltinEvalContext) InitProvider(ctx context.Context, addr addrs.AbsPro
 		return nil, err
 	}
 
-	if c.Evaluator != nil && c.Evaluator.Config != nil && c.Evaluator.Config.Module != nil {
-		// If an aliased provider is mocked, we use providerForTest wrapper.
-		// We cannot wrap providers.Factory itself, because factories don't support aliases.
-		pc, ok := c.Evaluator.Config.Module.GetProviderConfig(addr.Provider.Type, addr.Alias)
-		if ok && pc.IsMocked {
-			testP, err := newProviderForTestWithSchema(p, p.GetProviderSchema(ctx))
-			if err != nil {
-				return nil, err
-			}
-
-			p = testP.
-				withMockResources(pc.MockResources).
-				withOverrideResources(pc.OverrideResources)
-		}
-	}
-
 	log.Printf("[TRACE] BuiltinEvalContext: Initialized %q%s provider for %s", addr.String(), providerInstanceKey, addr)
 	c.ProviderCache[providerAddrKey][providerInstanceKey] = p
 

internal/tofu/node_provider.go

@@ -134,6 +134,11 @@ func (n *NodeApplyableProvider) ValidateProvider(ctx context.Context, evalCtx Ev
 	)
 	defer span.End()
 
+	if n.Config != nil && n.Config.IsMocked {
+		// Mocked for testing
+		return nil
+	}
+
 	configBody := buildProviderConfig(ctx, evalCtx, n.Addr, n.ProviderConfig())
 
 	// if a provider config is empty (only an alias), return early and don't continue
@@ -200,6 +205,11 @@ func (n *NodeApplyableProvider) ConfigureProvider(ctx context.Context, evalCtx E
 	)
 	defer span.End()
 
+	if n.Config != nil && n.Config.IsMocked {
+		// Mocked for testing
+		return nil
+	}
+
 	config := n.ProviderConfig()
 
 	configBody := buildProviderConfig(ctx, evalCtx, n.Addr, config)

internal/tofu/node_provider_abstract.go

@@ -88,6 +88,13 @@ func (n *NodeAbstractProvider) AttachProviderConfigSchema(schema *configschema.B
 	n.Schema = schema
 }
 
+func (n *NodeAbstractProvider) MocksAndOverrides() (IsMocked bool, MockResources []*configs.MockResource, OverrideResources []*configs.OverrideResource) {
+	if n.Config == nil {
+		return false, nil, nil
+	}
+	return n.Config.IsMocked, n.Config.MockResources, n.Config.OverrideResources
+}
+
 // GraphNodeDotter impl.
 func (n *NodeAbstractProvider) DotNode(name string, opts *dag.DotOpts) *dag.DotNode {
 	return &dag.DotNode{

internal/tofu/node_resource_abstract.go

@@ -584,7 +584,7 @@ func isResourceMovedToDifferentType(newAddr, oldAddr addrs.AbsResourceInstance)
 // the state.
 func (n *NodeAbstractResourceInstance) readResourceInstanceState(ctx context.Context, evalCtx EvalContext, addr addrs.AbsResourceInstance) (*states.ResourceInstanceObject, tfdiags.Diagnostics) {
 	var diags tfdiags.Diagnostics
-	provider, providerSchema, err := getProvider(ctx, evalCtx, n.ResolvedProvider.ProviderConfig, n.ResolvedProviderKey)
+	provider, providerSchema, err := n.getProvider(ctx, evalCtx)
 	if err != nil {
 		return nil, diags.Append(err)
 	}
@@ -639,7 +639,7 @@ func (n *NodeAbstractResourceInstance) readResourceInstanceState(ctx context.Con
 // instance in the state.
 func (n *NodeAbstractResourceInstance) readResourceInstanceStateDeposed(ctx context.Context, evalCtx EvalContext, addr addrs.AbsResourceInstance, key states.DeposedKey) (*states.ResourceInstanceObject, tfdiags.Diagnostics) {
 	var diags tfdiags.Diagnostics
-	provider, providerSchema, err := getProvider(ctx, evalCtx, n.ResolvedProvider.ProviderConfig, n.ResolvedProviderKey)
+	provider, providerSchema, err := n.getProvider(ctx, evalCtx)
 	if err != nil {
 		diags = diags.Append(err)
 		return nil, diags

internal/tofu/node_resource_abstract_instance.go

@@ -3220,24 +3220,41 @@ func (n *NodeAbstractResourceInstance) getProvider(ctx context.Context, evalCtx
 		return nil, providers.ProviderSchema{}, err
 	}
 
-	if n.Config == nil || !n.Config.IsOverridden {
-		if p, ok := underlyingProvider.(providerForTest); ok {
-			underlyingProvider = p.linkWithCurrentResource(n.Addr.ConfigResource())
+	var isOverridden bool
+	var overrideValues map[string]cty.Value
+
+	if n.ResolvedProvider.IsMocked {
+		isOverridden = true
+
+		// Mocked by the provider
+		for _, res := range n.ResolvedProvider.MockResources {
+			if res.Type == n.Addr.Resource.Resource.Type && res.Mode == n.Addr.Resource.Resource.Mode {
+				overrideValues = res.Defaults
+				break
+			}
 		}
 
-		return underlyingProvider, schema, nil
+		// Overridden by the provider (overrides mocks)
+		for _, res := range n.ResolvedProvider.OverrideResources {
+			if res.TargetParsed.Equal(n.Addr.ConfigResource()) && res.Mode == n.Addr.Resource.Resource.Mode {
+				overrideValues = res.Values
+				break
+			}
+		}
 	}
 
-	provider, err := newProviderForTestWithSchema(underlyingProvider, schema)
-	if err != nil {
-		return nil, providers.ProviderSchema{}, err
+	if n.Config != nil && n.Config.IsOverridden {
+		// Overridden in the currently running test (overrides any provider settings)
+		isOverridden = n.Config.IsOverridden
+		overrideValues = n.Config.OverrideValues
 	}
 
-	provider = provider.
-		withOverrideResource(n.Addr.ConfigResource(), n.Config.OverrideValues).
-		linkWithCurrentResource(n.Addr.ConfigResource())
+	if isOverridden {
+		provider, err := newProviderForTestWithSchema(underlyingProvider, schema, overrideValues)
+		return provider, schema, err
+	}
 
-	return provider, schema, nil
+	return underlyingProvider, schema, err
 }
 
 func maybeImproveResourceInstanceDiagnostics(diags tfdiags.Diagnostics, excludeAddr addrs.Targetable) tfdiags.Diagnostics {

internal/tofu/node_resource_apply_instance.go

@@ -212,7 +212,7 @@ func (n *NodeApplyableResourceInstance) ephemeralResourceExecute(ctx context.Con
 }
 
 func (n *NodeApplyableResourceInstance) dataResourceExecute(ctx context.Context, evalCtx EvalContext) (diags tfdiags.Diagnostics) {
-	_, providerSchema, err := getProvider(ctx, evalCtx, n.ResolvedProvider.ProviderConfig, n.ResolvedProviderKey)
+	_, providerSchema, err := n.getProvider(ctx, evalCtx)
 	diags = diags.Append(err)
 	if diags.HasErrors() {
 		return diags
@@ -281,7 +281,7 @@ func (n *NodeApplyableResourceInstance) managedResourceExecute(ctx context.Conte
 	var deposedKey states.DeposedKey
 
 	addr := n.ResourceInstanceAddr().Resource
-	_, providerSchema, err := getProvider(ctx, evalCtx, n.ResolvedProvider.ProviderConfig, n.ResolvedProviderKey)
+	_, providerSchema, err := n.getProvider(ctx, evalCtx)
 	diags = diags.Append(err)
 	if diags.HasErrors() {
 		return diags

internal/tofu/node_resource_deposed.go

@@ -386,7 +386,7 @@ func (n *NodeDestroyDeposedResourceInstanceObject) writeResourceInstanceState(ct
 		return nil
 	}
 
-	_, providerSchema, err := getProvider(ctx, evalCtx, n.ResolvedProvider.ProviderConfig, n.ResolvedProviderKey)
+	_, providerSchema, err := n.getProvider(ctx, evalCtx)
 	if err != nil {
 		return err
 	}

internal/tofu/node_resource_destroy.go

@@ -203,7 +203,7 @@ func (n *NodeDestroyResourceInstance) managedResourceExecute(ctx context.Context
 	var changeApply *plans.ResourceInstanceChange
 	var state *states.ResourceInstanceObject
 
-	_, providerSchema, err := getProvider(ctx, evalCtx, n.ResolvedProvider.ProviderConfig, n.ResolvedProviderKey)
+	_, providerSchema, err := n.getProvider(ctx, evalCtx)
 	diags = diags.Append(err)
 	if diags.HasErrors() {
 		return diags

internal/tofu/node_resource_plan_instance.go

@@ -136,7 +136,7 @@ func (n *NodePlannableResourceInstance) dataResourceExecute(ctx context.Context,
 
 	var change *plans.ResourceInstanceChange
 
-	_, providerSchema, err := getProvider(ctx, evalCtx, n.ResolvedProvider.ProviderConfig, n.ResolvedProviderKey)
+	_, providerSchema, err := n.getProvider(ctx, evalCtx)
 	diags = diags.Append(err)
 	if diags.HasErrors() {
 		return diags
@@ -191,7 +191,7 @@ func (n *NodePlannableResourceInstance) ephemeralResourceExecute(ctx context.Con
 	config := n.Config
 	addr := n.ResourceInstanceAddr()
 
-	_, providerSchema, err := getProvider(ctx, evalCtx, n.ResolvedProvider.ProviderConfig, n.ResolvedProviderKey)
+	_, providerSchema, err := n.getProvider(ctx, evalCtx)
 	diags = diags.Append(err)
 	if diags.HasErrors() {
 		return diags
@@ -251,7 +251,7 @@ func (n *NodePlannableResourceInstance) managedResourceExecute(ctx context.Conte
 		checkRuleSeverity = tfdiags.Warning
 	}
 
-	provider, providerSchema, err := getProvider(ctx, evalCtx, n.ResolvedProvider.ProviderConfig, n.ResolvedProviderKey)
+	provider, providerSchema, err := n.getProvider(ctx, evalCtx)
 	diags = diags.Append(err)
 	if diags.HasErrors() {
 		return diags

internal/tofu/provider_for_test_framework.go

@@ -11,7 +11,6 @@ import (
 	"hash/fnv"
 
 	"github.com/opentofu/opentofu/internal/addrs"
-	"github.com/opentofu/opentofu/internal/configs"
 	"github.com/opentofu/opentofu/internal/configs/hcl2shim"
 	"github.com/opentofu/opentofu/internal/providers"
 	"github.com/opentofu/opentofu/internal/tfdiags"
@@ -20,43 +19,28 @@ import (
 
 var _ providers.Interface = &providerForTest{}
 
-// providerForTest is a wrapper around a real provider to allow certain resources to be overridden
-// (by address) or mocked (by provider and resource type) for testing framework.
+// providerForTest is a wrapper around a real provider to allow a specific resource to be overridden
+// for the testing framework.  It is used by [NodeResourceAbstractInstance.getProvider] to fulfil the mocks
+// and overrides for that one specific resource instance, any other usage is a bug in OpenTofu and should
+// be corrected.
 type providerForTest struct {
 	// providers.Interface is not embedded to make it safer to extend
 	// the interface without silently breaking providerForTest functionality.
 	internal providers.Interface
 	schema   providers.ProviderSchema
 
-	mockResources     mockResourcesForTest
-	overrideResources overrideResourcesForTest
-
-	currentResourceAddress string
+	overrideValues map[string]cty.Value
 }
 
-func newProviderForTestWithSchema(internal providers.Interface, schema providers.ProviderSchema) (providerForTest, error) {
-	if p, ok := internal.(providerForTest); ok {
-		// We can create a proper deep copy here, however currently
-		// it is only relevant for override resources, since we extend
-		// the override resource map in NodeAbstractResourceInstance.
-		return p.withCopiedOverrideResources(), nil
-	}
-
+func newProviderForTestWithSchema(internal providers.Interface, schema providers.ProviderSchema, overrideValues map[string]cty.Value) (providerForTest, error) {
 	if schema.Diagnostics.HasErrors() {
 		return providerForTest{}, fmt.Errorf("invalid provider schema for test wrapper: %w", schema.Diagnostics.Err())
 	}
 
 	return providerForTest{
-		internal: internal,
-		schema:   schema,
-		mockResources: mockResourcesForTest{
-			managed: make(map[string]resourceForTest),
-			data:    make(map[string]resourceForTest),
-		},
-		overrideResources: overrideResourcesForTest{
-			managed: make(map[string]resourceForTest),
-			data:    make(map[string]resourceForTest),
-		},
+		internal:       internal,
+		schema:         schema,
+		overrideValues: overrideValues,
 	}, nil
 }
 
@@ -84,12 +68,10 @@ func (p providerForTest) PlanResourceChange(_ context.Context, r providers.PlanR
 
 	resSchema, _ := p.schema.SchemaForResourceType(addrs.ManagedResourceMode, r.TypeName)
 
-	mockValues := p.getMockValuesForManagedResource(r.TypeName)
-
 	var resp providers.PlanResourceChangeResponse
 
 	resp.PlannedState, resp.Diagnostics = newMockValueComposer(r.TypeName).
-		ComposeBySchema(resSchema, r.Config, mockValues)
+		ComposeBySchema(resSchema, r.Config, p.overrideValues)
 
 	return resp
 }
@@ -105,10 +87,8 @@ func (p providerForTest) ReadDataSource(_ context.Context, r providers.ReadDataS
 
 	var resp providers.ReadDataSourceResponse
 
-	mockValues := p.getMockValuesForDataResource(r.TypeName)
-
 	resp.State, resp.Diagnostics = newMockValueComposer(r.TypeName).
-		ComposeBySchema(resSchema, r.Config, mockValues)
+		ComposeBySchema(resSchema, r.Config, p.overrideValues)
 
 	return resp
 }
@@ -195,125 +175,6 @@ func (p providerForTest) Close(ctx context.Context) error {
 	return p.internal.Close(ctx)
 }
 
-func (p providerForTest) withMockResources(mockResources []*configs.MockResource) providerForTest {
-	for _, res := range mockResources {
-		var resources map[mockResourceType]resourceForTest
-
-		switch res.Mode {
-		case addrs.ManagedResourceMode:
-			resources = p.mockResources.managed
-		case addrs.DataResourceMode:
-			resources = p.mockResources.data
-		case addrs.InvalidResourceMode:
-			panic("BUG: invalid mock resource mode")
-		default:
-			panic("BUG: unsupported mock resource mode: " + res.Mode.String())
-		}
-
-		resources[res.Type] = resourceForTest{
-			values: res.Defaults,
-		}
-	}
-
-	return p
-}
-
-func (p providerForTest) withCopiedOverrideResources() providerForTest {
-	p.overrideResources = p.overrideResources.copy()
-	return p
-}
-
-func (p providerForTest) withOverrideResources(overrideResources []*configs.OverrideResource) providerForTest {
-	for _, res := range overrideResources {
-		p = p.withOverrideResource(*res.TargetParsed, res.Values)
-	}
-
-	return p
-}
-
-func (p providerForTest) withOverrideResource(addr addrs.ConfigResource, overrides map[string]cty.Value) providerForTest {
-	var resources map[string]resourceForTest
-
-	switch addr.Resource.Mode {
-	case addrs.ManagedResourceMode:
-		resources = p.overrideResources.managed
-	case addrs.DataResourceMode:
-		resources = p.overrideResources.data
-	case addrs.InvalidResourceMode:
-		panic("BUG: invalid override resource mode")
-	default:
-		panic("BUG: unsupported override resource mode: " + addr.Resource.Mode.String())
-	}
-
-	resources[addr.String()] = resourceForTest{
-		values: overrides,
-	}
-
-	return p
-}
-
-func (p providerForTest) linkWithCurrentResource(addr addrs.ConfigResource) providerForTest {
-	p.currentResourceAddress = addr.String()
-	return p
-}
-
-type resourceForTest struct {
-	values map[string]cty.Value
-}
-
-type mockResourceType = string
-
-type mockResourcesForTest struct {
-	managed map[mockResourceType]resourceForTest
-	data    map[mockResourceType]resourceForTest
-}
-
-type overrideResourceAddress = string
-
-type overrideResourcesForTest struct {
-	managed map[overrideResourceAddress]resourceForTest
-	data    map[overrideResourceAddress]resourceForTest
-}
-
-func (res overrideResourcesForTest) copy() overrideResourcesForTest {
-	resCopy := overrideResourcesForTest{
-		managed: make(map[overrideResourceAddress]resourceForTest, len(res.managed)),
-		data:    make(map[overrideResourceAddress]resourceForTest, len(res.data)),
-	}
-
-	for k, v := range res.managed {
-		resCopy.managed[k] = v
-	}
-
-	for k, v := range res.data {
-		resCopy.data[k] = v
-	}
-
-	return resCopy
-}
-
-func (p providerForTest) getMockValuesForManagedResource(typeName string) map[string]cty.Value {
-	if p.currentResourceAddress != "" {
-		res, ok := p.overrideResources.managed[p.currentResourceAddress]
-		if ok {
-			return res.values
-		}
-	}
-
-	return p.mockResources.managed[typeName].values
-}
-
-func (p providerForTest) getMockValuesForDataResource(typeName string) map[string]cty.Value {
-	if p.currentResourceAddress != "" {
-		res, ok := p.overrideResources.data[p.currentResourceAddress]
-		if ok {
-			return res.values
-		}
-	}
-
-	return p.mockResources.data[typeName].values
-}
-
 func newMockValueComposer(typeName string) hcl2shim.MockValueComposer {
 	hash := fnv.New32()
 	hash.Write([]byte(typeName))

internal/tofu/provider_for_test_framework_test.go

@@ -15,7 +15,7 @@ import (
 func TestProviderForTest_ReadResource(t *testing.T) {
 	mockProvider := &MockProvider{}
 
-	provider, err := newProviderForTestWithSchema(mockProvider, mockProvider.GetProviderSchema(t.Context()))
+	provider, err := newProviderForTestWithSchema(mockProvider, mockProvider.GetProviderSchema(t.Context()), nil)
 	if err != nil {
 		t.Fatalf("unexpected error: %s", err.Error())
 	}

internal/tofu/transform_provider.go

@@ -53,6 +53,8 @@ type GraphNodeProvider interface {
 	GraphNodeModulePath
 	ProviderAddr() addrs.AbsProviderConfig
 	Name() string
+	// For test framework
+	MocksAndOverrides() (IsMocked bool, MockResources []*configs.MockResource, OverrideResources []*configs.OverrideResource)
 }
 
 // GraphNodeCloseProvider is an interface that nodes that can be a close
@@ -77,6 +79,11 @@ type ResolvedProvider struct {
 	KeyModule      addrs.Module
 	KeyResource    bool
 	KeyExact       addrs.InstanceKey
+
+	// Test overrides
+	IsMocked          bool
+	MockResources     []*configs.MockResource
+	OverrideResources []*configs.OverrideResource
 }
 
 // GraphNodeProviderConsumer is an interface that nodes that require
@@ -169,14 +176,18 @@ func (t *ProviderTransformer) Transform(_ context.Context, g *Graph) error {
 			}
 
 			log.Printf("[DEBUG] ProviderTransformer: %q (%T) needs exactly %s", dag.VertexName(v), v, dag.VertexName(target))
-			pv.SetProvider(ResolvedProvider{
+
+			resolved := ResolvedProvider{
 				ProviderConfig: target.ProviderAddr(),
 				// Pass through key data
 				KeyExpression: req.KeyExpression,
 				KeyModule:     req.KeyModule,
 				KeyResource:   req.KeyResource,
 				KeyExact:      req.KeyExact,
-			})
+			}
+			resolved.IsMocked, resolved.MockResources, resolved.OverrideResources = target.MocksAndOverrides()
+			pv.SetProvider(resolved)
+
 			g.Connect(dag.BasicEdge(v, target))
 		case addrs.LocalProviderConfig:
 			// We assume that the value returned from Provider() has already been
@@ -244,6 +255,9 @@ func (t *ProviderTransformer) Transform(_ context.Context, g *Graph) error {
 			}
 			resolved.ProviderConfig = target.ProviderAddr()
 
+			// Include test mocking and override extensions
+			resolved.IsMocked, resolved.MockResources, resolved.OverrideResources = target.MocksAndOverrides()
+
 			log.Printf("[DEBUG] ProviderTransformer: %q (%T) needs %s", dag.VertexName(v), v, dag.VertexName(target))
 			pv.SetProvider(resolved)
 			g.Connect(dag.BasicEdge(v, target))
@@ -693,6 +707,10 @@ func (n *graphNodeProxyProvider) Target() GraphNodeProvider {
 	}
 }
 
+func (n *graphNodeProxyProvider) MocksAndOverrides() (IsMocked bool, MockResources []*configs.MockResource, OverrideResources []*configs.OverrideResource) {
+	return n.Target().MocksAndOverrides()
+}
+
 // Find the *single* keyExpression that is used in the provider
 // chain.  This is not ideal, but it works with current constraints on this feature
 func (n *graphNodeProxyProvider) TargetExpr() (hcl.Expression, addrs.Module) {