Repository: pulumi/pulumi-terraform-bridge
Pull request #997 (2023-04-15 00:24:37)

Change Muxer API to take a single tfbridge.ProviderInfo
This builds on https://github.com/pulumi/pulumi-terraform-bridge/pull/996, simplify the user facing API.

Example usage looks like https://github.com/pulumi/pulumi-cloudflare/pull/379.


Patch:
diff --git a/pf/tfgen/extend.go b/pf/internal/muxer/muxer.go
similarity index 87%
rename from pf/tfgen/extend.go
rename to pf/internal/muxer/muxer.go
index e2ef24cb..af89dae7 100644
--- a/pf/tfgen/extend.go
+++ b/pf/internal/muxer/muxer.go
@@ -12,7 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-package tfgen
+package muxer
 
 import (
 	"fmt"
@@ -35,10 +35,10 @@ func SchemaOnlyPluginFrameworkProvider(ctx context.Context, provider provider.Pr
 }
 
 // AugmentShimWithPF augments an existing shim with a PF provider.Provider.
-func AugmentShimWithPF(ctx context.Context, shim shim.Provider, pf provider.Provider) shim.Provider {
+func AugmentShimWithPF(ctx context.Context, shim shim.Provider, pf provider.Provider) *ProviderShim {
 
-	var p mergedProviderShim
-	if alreadyMerged, ok := shim.(*mergedProviderShim); ok {
+	var p ProviderShim
+	if alreadyMerged, ok := shim.(*ProviderShim); ok {
 		p = *alreadyMerged
 	} else {
 		p = newMergedProviderShim(shim)
@@ -50,13 +50,13 @@ func AugmentShimWithPF(ctx context.Context, shim shim.Provider, pf provider.Prov
 }
 
 // A merged `shim.Provider` that remembers which `shim.Provider`s it is composed of.
-type mergedProviderShim struct {
+type ProviderShim struct {
 	simpleSchemaProvider
 
-	providers []shim.Provider
+	MuxedProviders []shim.Provider
 }
 
-func (m *mergedProviderShim) extend(provider shim.Provider) error {
+func (m *ProviderShim) extend(provider shim.Provider) error {
 	res, err := disjointUnion(m.resources, provider.ResourcesMap())
 	if err != nil {
 		return err
@@ -69,28 +69,28 @@ func (m *mergedProviderShim) extend(provider shim.Provider) error {
 
 	m.resources = res
 	m.dataSources = data
-	m.providers = append(m.providers, provider)
+	m.MuxedProviders = append(m.MuxedProviders, provider)
 	return nil
 }
 
-func newMergedProviderShim(provider shim.Provider) mergedProviderShim {
-	return mergedProviderShim{
+func newMergedProviderShim(provider shim.Provider) ProviderShim {
+	return ProviderShim{
 		simpleSchemaProvider: simpleSchemaProvider{
 			schema:      provider.Schema(),
 			resources:   provider.ResourcesMap(),
 			dataSources: provider.DataSourcesMap(),
 		},
-		providers: []shim.Provider{provider},
+		MuxedProviders: []shim.Provider{provider},
 	}
 }
 
-func (m *mergedProviderShim) resolveDispatch(info *tfbridge.ProviderInfo) muxer.DispatchTable {
+func (m *ProviderShim) ResolveDispatch(info *tfbridge.ProviderInfo) muxer.DispatchTable {
 	var dispatch muxer.DispatchTable
 	dispatch.Resources = map[string]int{}
 	dispatch.Functions = map[string]int{}
 
 	for tfToken, res := range info.Resources {
-		for i, p := range m.providers {
+		for i, p := range m.MuxedProviders {
 			_, ok := p.ResourcesMap().GetOk(tfToken)
 			if ok {
 				dispatch.Resources[string(res.GetTok())] = i
@@ -99,7 +99,7 @@ func (m *mergedProviderShim) resolveDispatch(info *tfbridge.ProviderInfo) muxer.
 		}
 	}
 	for tfToken, data := range info.DataSources {
-		for i, p := range m.providers {
+		for i, p := range m.MuxedProviders {
 			_, ok := p.DataSourcesMap().GetOk(tfToken)
 			if ok {
 				dispatch.Functions[string(data.GetTok())] = i
diff --git a/pf/internal/schemashim/provider.go b/pf/internal/schemashim/provider.go
index 8c69148e..c1605871 100644
--- a/pf/internal/schemashim/provider.go
+++ b/pf/internal/schemashim/provider.go
@@ -27,6 +27,10 @@ type SchemaOnlyProvider struct {
 	tf  pfprovider.Provider
 }
 
+func (p *SchemaOnlyProvider) PfProvider() pfprovider.Provider {
+	return p.tf
+}
+
 var _ shim.Provider = (*SchemaOnlyProvider)(nil)
 
 func (p *SchemaOnlyProvider) Schema() shim.SchemaMap {
diff --git a/pf/tfbridge/extend.go b/pf/tfbridge/extend.go
new file mode 100644
index 00000000..b803af5d
--- /dev/null
+++ b/pf/tfbridge/extend.go
@@ -0,0 +1,33 @@
+// Copyright 2016-2023, Pulumi Corporation.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tfbridge
+
+import (
+	"context"
+
+	pfprovider "github.com/hashicorp/terraform-plugin-framework/provider"
+	"github.com/pulumi/pulumi-terraform-bridge/pf/internal/muxer"
+	"github.com/pulumi/pulumi-terraform-bridge/pf/internal/schemashim"
+	shim "github.com/pulumi/pulumi-terraform-bridge/v3/pkg/tfshim"
+)
+
+func SchemaOnlyPluginFrameworkProvider(ctx context.Context, p pfprovider.Provider) shim.Provider {
+	return schemashim.ShimSchemaOnlyProvider(ctx, p)
+}
+
+// AugmentShimWithPF augments an existing shim with a PF provider.Provider.
+func AugmentShimWithPF(ctx context.Context, shim shim.Provider, p pfprovider.Provider) shim.Provider {
+	return muxer.AugmentShimWithPF(ctx, shim, p)
+}
diff --git a/pf/tfbridge/main.go b/pf/tfbridge/main.go
index ae17a6c2..b4b2ae92 100644
--- a/pf/tfbridge/main.go
+++ b/pf/tfbridge/main.go
@@ -27,8 +27,9 @@ import (
 	"github.com/pulumi/pulumi/sdk/v3/go/common/util/contract"
 	pulumirpc "github.com/pulumi/pulumi/sdk/v3/proto/go"
 
+	pfmuxer "github.com/pulumi/pulumi-terraform-bridge/pf/internal/muxer"
+	"github.com/pulumi/pulumi-terraform-bridge/pf/internal/schemashim"
 	"github.com/pulumi/pulumi-terraform-bridge/v3/pkg/tfbridge"
-	sdkBridge "github.com/pulumi/pulumi-terraform-bridge/v3/pkg/tfbridge"
 	"github.com/pulumi/pulumi-terraform-bridge/v3/unstable/metadata"
 	"github.com/pulumi/pulumi-terraform-bridge/x/muxer"
 )
@@ -91,10 +92,10 @@ func handleFlags(ctx context.Context, prov ProviderInfo, meta ProviderMetadata,
 // Implements main() or a bridged Pulumi plugin, complete with argument parsing.
 //
 // This is an experimental API.
-func MainWithMuxer(ctx context.Context, pkg string, meta ProviderMetadata, infos ...Muxed) {
-	f := MakeMuxedServer(ctx, pkg, meta, infos...)
+func MainWithMuxer(ctx context.Context, schema []byte, info tfbridge.ProviderInfo) {
+	f := MakeMuxedServer(ctx, schema, info)
 
-	err := rprovider.Main(pkg, f)
+	err := rprovider.Main(info.Name, f)
 	if err != nil {
 		cmdutil.ExitError(err.Error())
 	}
@@ -105,11 +106,13 @@ func MainWithMuxer(ctx context.Context, pkg string, meta ProviderMetadata, infos
 // This function exposes implementation details for testing. It should not be used outside
 // of pulumi-terraform-bridge.  This is an experimental API.
 func MakeMuxedServer(
-	ctx context.Context, pkg string, meta ProviderMetadata, infos ...Muxed,
+	ctx context.Context, schema []byte, info tfbridge.ProviderInfo,
 ) func(host *rprovider.HostClient) (pulumirpc.ResourceProviderServer, error) {
-	version := infos[0].GetInfo().Version
-	schema := string(meta.PackageSchema)
-	dispatchTable, found, err := metadata.Get[muxer.DispatchTable](infos[0].GetInfo().GetMetadata(), "mux")
+
+	shim, ok := info.P.(*pfmuxer.ProviderShim)
+	contract.Assertf(ok, "MainWithMuxer must have a ProviderInfo.P created with AugmentShimWithPF")
+	version := info.Version
+	dispatchTable, found, err := metadata.Get[muxer.DispatchTable](info.GetMetadata(), "mux")
 	if err != nil {
 		cmdutil.ExitError(err.Error())
 	}
@@ -119,30 +122,33 @@ func MakeMuxedServer(
 	}
 	m := muxer.Main{
 		DispatchTable: dispatchTable,
-		Schema:        schema,
+		Schema:        string(schema),
 		GetMappingHandler: map[string]muxer.MultiMappingHandler{
 			"tf":        combineTFGetMappingKey,
 			"terraform": combineTFGetMappingKey,
 		},
 	}
 	return func(host *rprovider.HostClient) (pulumirpc.ResourceProviderServer, error) {
-		for _, info := range infos {
-			info := info // https://github.com/golang/go/wiki/CommonMistakes#using-goroutines-on-loop-iterator-variables
+		for _, prov := range shim.MuxedProviders {
+			prov := prov // https://github.com/golang/go/wiki/CommonMistakes#using-goroutines-on-loop-iterator-variables
 
-			// Add PF based servers to the runtime.
-			if info.PF != nil {
+			switch prov := prov.(type) {
+			case *schemashim.SchemaOnlyProvider:
+				m.Servers = append(m.Servers, muxer.Endpoint{
+					Server: func(host *rprovider.HostClient) (pulumirpc.ResourceProviderServer, error) {
+						return newProviderServer(ctx, host, ProviderInfo{
+							ProviderInfo: info,
+							NewProvider:  prov.PfProvider,
+						}, ProviderMetadata{PackageSchema: schema})
+					}})
+			default:
 				m.Servers = append(m.Servers, muxer.Endpoint{
 					Server: func(host *rprovider.HostClient) (pulumirpc.ResourceProviderServer, error) {
-						return newProviderServer(ctx, host, *info.PF, meta)
+						return tfbridge.NewProvider(ctx, host, info.Name, version, prov, info, schema), nil
 					}})
-				continue
 			}
-			m.Servers = append(m.Servers, muxer.Endpoint{
-				Server: func(host *rprovider.HostClient) (pulumirpc.ResourceProviderServer, error) {
-					return tfbridge.NewProvider(ctx, host, pkg, version, info.SDK.P, *info.SDK, meta.PackageSchema), nil
-				}})
 		}
-		return m.Server(host, pkg, version)
+		return m.Server(host, info.Name, version)
 	}
 }
 
@@ -212,19 +218,3 @@ func mergeMapUnder[K comparable, V any](dst *map[K]V, src map[K]V) {
 		}
 	}
 }
-
-// A union of pf and sdk based ProviderInfo for use in MainWithMuxer.
-// Exactly 1 field of this struct should hold a value
-//
-// This is an experimental API.
-type Muxed struct {
-	PF  *ProviderInfo
-	SDK *sdkBridge.ProviderInfo
-}
-
-func (m Muxed) GetInfo() sdkBridge.ProviderInfo {
-	if m.PF == nil {
-		return *m.SDK
-	}
-	return m.PF.ProviderInfo
-}
diff --git a/pf/tfgen/main.go b/pf/tfgen/main.go
index 3bbdf247..73d7af37 100644
--- a/pf/tfgen/main.go
+++ b/pf/tfgen/main.go
@@ -18,16 +18,13 @@ import (
 	"context"
 	"fmt"
 
-	"github.com/hashicorp/go-multierror"
-	"github.com/pulumi/pulumi/pkg/v3/codegen/schema"
-	"github.com/pulumi/pulumi/sdk/v3/go/common/util/cmdutil"
 	"github.com/pulumi/pulumi/sdk/v3/go/common/util/contract"
 
+	pfmuxer "github.com/pulumi/pulumi-terraform-bridge/pf/internal/muxer"
 	"github.com/pulumi/pulumi-terraform-bridge/pf/tfbridge"
 	sdkBridge "github.com/pulumi/pulumi-terraform-bridge/v3/pkg/tfbridge"
 	"github.com/pulumi/pulumi-terraform-bridge/v3/pkg/tfgen"
 	"github.com/pulumi/pulumi-terraform-bridge/v3/unstable/metadata"
-	"github.com/pulumi/pulumi-terraform-bridge/x/muxer"
 )
 
 // Implements main() logic for a provider build-time helper utility. By convention these utilities are named
@@ -76,29 +73,17 @@ func Main(provider string, info tfbridge.ProviderInfo) {
 // This is an experimental API.
 //
 // [Pulumi Package Schema]: https://www.pulumi.com/docs/guides/pulumi-packages/schema/
-func MainWithMuxer(provider string, infos ...tfbridge.Muxed) {
-	ctx := context.Background()
-
-	contract.Assertf(len(infos) == 2,
-		"We expect to see a SDKv2 and then a provider info for a muxed provider")
-	sdkV2 := infos[0].SDK
-	pf := *infos[1].PF
-
-	shim := AugmentShimWithPF(ctx, sdkV2.P, pf.NewProvider()).(*mergedProviderShim)
+func MainWithMuxer(info sdkBridge.ProviderInfo) {
+	shim, ok := info.P.(*pfmuxer.ProviderShim)
+	contract.Assertf(ok, "MainWithMuxer must have a ProviderInfo.P created with AugmentShimWithPF")
 
-	version := sdkV2.Version
-	sdkV2.P = shim
+	tfgen.MainWithCustomGenerate(info.Name, info.Version, info, func(opts tfgen.GeneratorOptions) error {
 
-	mergeDisjoint(&sdkV2.Resources, pf.Resources)
-	mergeDisjoint(&sdkV2.DataSources, pf.DataSources)
-
-	tfgen.MainWithCustomGenerate(provider, version, *sdkV2, func(opts tfgen.GeneratorOptions) error {
-
-		if sdkV2.MetadataInfo == nil {
+		if info.MetadataInfo == nil {
 			return fmt.Errorf("ProviderInfo.MetadataInfo is required and cannot be nil")
 		}
 
-		err := metadata.Set(sdkV2.GetMetadata(), "mux", shim.resolveDispatch(sdkV2))
+		err := metadata.Set(info.GetMetadata(), "mux", shim.ResolveDispatch(&info))
 		if err != nil {
 			return err
 		}
@@ -121,181 +106,3 @@ func MainWithMuxer(provider string, infos ...tfbridge.Muxed) {
 		return nil
 	})
 }
-
-func mergeDisjoint[T any](dst *map[string]T, src map[string]T) {
-	contract.Assertf(dst != nil, "Cannot assign to a nil destination")
-	if src == nil {
-		return
-	}
-	if *dst == nil {
-		*dst = map[string]T{}
-	}
-	for k, v := range src {
-		_, ok := (*dst)[k]
-		contract.Assertf(!ok, "Map was not disjoint on key %s", k)
-		(*dst)[k] = v
-	}
-}
-
-// Mux multiple infos into a unified representation.
-//
-// This function exposes implementation details for testing. It should not be used outside
-// of pulumi-terraform-bridge.  This is an experimental API.
-func UnstableMuxProviderInfo(
-	ctx context.Context, opts tfgen.GeneratorOptions,
-	provider string, infos ...tfbridge.Muxed,
-) (sdkBridge.ProviderInfo, muxer.DispatchTable, schema.PackageSpec, tfgen.Renames, error) {
-	if opts.Sink == nil {
-		opts.Sink = cmdutil.Diag()
-	}
-	muxedInfo := infos[0].GetInfo()
-	schemas := []schema.PackageSpec{}
-	var errs multierror.Error
-
-	var renames []tfgen.Renames
-
-	for i, union := range infos {
-		i, union := i, union
-
-		anErr := func(err error) bool {
-			if err == nil {
-				return false
-			}
-			errs.Errors = append(errs.Errors, fmt.Errorf("Muxed proider[%d]: %w", i, err))
-			return true
-		}
-
-		info := union.SDK
-		if pfInfo := union.PF; pfInfo != nil {
-			if err := notSupported(opts.Sink, pfInfo.ProviderInfo); anErr(err) {
-				continue
-			}
-
-			shimInfo := shimSchemaOnlyProviderInfo(ctx, *pfInfo)
-			info = &shimInfo
-		}
-
-		// Initially, muxedInfo **is** infos[0]. We don't want to perform
-		// any destructive edits on already defined functions.
-		if i != 0 {
-			// Now we purge already defined items from the next schema,
-			// marking the all remaining items as defined for future schemas.
-
-			muxedInfo.IgnoreMappings = append(muxedInfo.IgnoreMappings,
-				info.IgnoreMappings...)
-
-			// Resources
-			layerUnder(&muxedInfo.Resources, info.Resources, muxedInfo.ExtraResources,
-				func(k string) {
-					delete(info.Resources, k)
-					info.IgnoreMappings = append(info.IgnoreMappings, k)
-				})
-
-			layerUnder(&muxedInfo.ExtraResources, info.ExtraResources, muxedInfo.Resources,
-				func(k string) {
-					delete(info.ExtraResources, k)
-					info.IgnoreMappings = append(info.IgnoreMappings, k)
-				})
-
-			// Datasources/functions
-			layerUnder(&muxedInfo.DataSources, info.DataSources, muxedInfo.ExtraFunctions,
-				func(k string) {
-					delete(info.DataSources, k)
-					info.IgnoreMappings = append(info.IgnoreMappings, k)
-				})
-			layerUnder(&muxedInfo.ExtraFunctions, info.ExtraFunctions, muxedInfo.DataSources,
-				func(k string) {
-					delete(info.ExtraFunctions, k)
-					info.IgnoreMappings = append(info.IgnoreMappings, k)
-				})
-
-			// Config
-			layerUnder(&muxedInfo.Config, info.Config, muxedInfo.ExtraConfig,
-				func(k string) { delete(info.Config, k) })
-			layerUnder(&muxedInfo.ExtraConfig, info.ExtraConfig, muxedInfo.Config,
-				func(k string) { delete(info.Config, k) })
-		}
-
-		genResult, err := tfgen.GenerateSchemaWithOptions(tfgen.GenerateSchemaOptions{
-			ProviderInfo:    *info,
-			DiagnosticsSink: opts.Sink,
-		})
-
-		if anErr(err) {
-			continue
-		}
-
-		renames = append(renames, genResult.Renames)
-		schemas = append(schemas, genResult.PackageSpec)
-	}
-
-	if err := errs.ErrorOrNil(); err != nil {
-		return sdkBridge.ProviderInfo{}, muxer.DispatchTable{},
-			schema.PackageSpec{}, tfgen.Renames{}, err
-	}
-
-	m, s, err := muxer.MergeSchemasAndComputeDispatchTable(schemas)
-	return muxedInfo, m, s, mergeRenames(renames), err
-}
-
-func mergeRenames(renames []tfgen.Renames) tfgen.Renames {
-	main := renames[0]
-	for _, rename := range renames[1:] {
-		for k, v := range rename.Resources {
-			_, exists := main.Resources[k]
-			if !exists {
-				main.Resources[k] = v
-			}
-		}
-
-		for k, v := range rename.Functions {
-			_, exists := main.Functions[k]
-			if !exists {
-				main.Functions[k] = v
-			}
-		}
-		for k, v := range rename.RenamedProperties {
-			_, exists := main.RenamedProperties[k]
-			if !exists {
-				main.RenamedProperties[k] = v
-			}
-		}
-		for k, v := range rename.RenamedConfigProperties {
-			_, exists := main.RenamedConfigProperties[k]
-			if !exists {
-				main.RenamedConfigProperties[k] = v
-			}
-		}
-		for k, v := range rename.Resources {
-			_, exists := main.Resources[k]
-			if !exists {
-				main.Resources[k] = v
-			}
-		}
-	}
-	return main
-}
-
-func layerUnder[T any, O any](
-	dst *map[string]T, src map[string]T, alternativeDst map[string]O,
-	afterAdd func(k string),
-) {
-	contract.Assertf(dst != nil, "Cannot assign to a nil destination")
-	if src == nil {
-		return
-	}
-	if *dst == nil {
-		*dst = map[string]T{}
-	}
-	for k, v := range src {
-		_, ok := (*dst)[k]
-		if alternativeDst != nil && !ok {
-			_, ok = alternativeDst[k]
-		}
-		if ok {
-			continue
-		}
-		(*dst)[k] = v
-		afterAdd(k)
-	}
-}


Test patch:
diff --git a/pf/tests/internal/testprovider/cmd/pulumi-resource-muxedrandom/main.go b/pf/tests/internal/testprovider/cmd/pulumi-resource-muxedrandom/main.go
index 7aa15bf0..b1ef8767 100644
--- a/pf/tests/internal/testprovider/cmd/pulumi-resource-muxedrandom/main.go
+++ b/pf/tests/internal/testprovider/cmd/pulumi-resource-muxedrandom/main.go
@@ -27,8 +27,6 @@ import (
 var schema []byte
 
 func main() {
-	meta := tfbridge.ProviderMetadata{PackageSchema: schema}
-	tfbridge.MainWithMuxer(context.Background(), "random", meta,
-		testprovider.MuxedRandomProvider()...,
-	)
+	tfbridge.MainWithMuxer(context.Background(), schema,
+		testprovider.MuxedRandomProvider())
 }
diff --git a/pf/tests/internal/testprovider/cmd/pulumi-tfgen-muxedrandom/main.go b/pf/tests/internal/testprovider/cmd/pulumi-tfgen-muxedrandom/main.go
index fdd09caf..24575c6a 100644
--- a/pf/tests/internal/testprovider/cmd/pulumi-tfgen-muxedrandom/main.go
+++ b/pf/tests/internal/testprovider/cmd/pulumi-tfgen-muxedrandom/main.go
@@ -22,7 +22,5 @@ import (
 )
 
 func main() {
-	tfgen.MainWithMuxer("random",
-		testprovider.MuxedRandomProvider()...,
-	)
+	tfgen.MainWithMuxer(testprovider.MuxedRandomProvider())
 }
diff --git a/pf/tests/internal/testprovider/random.go b/pf/tests/internal/testprovider/random.go
index 03184aef..c8b5bb73 100644
--- a/pf/tests/internal/testprovider/random.go
+++ b/pf/tests/internal/testprovider/random.go
@@ -15,6 +15,7 @@
 package testprovider
 
 import (
+	"context"
 	_ "embed"
 	"fmt"
 	"path/filepath"
@@ -139,8 +140,14 @@ func RandomSDKProvider() tfbridge.ProviderInfo {
 	}
 
 	info := tfbridge.ProviderInfo{
-		Name: "muxedrandom",
-		P:    sdkv2.NewProvider(sdkv2randomprovider.New()),
+		Name:        "muxedrandom",
+		Description: "A Pulumi package to safely use randomness in Pulumi programs.",
+		Keywords:    []string{"pulumi", "random"},
+		License:     "Apache-2.0",
+		Homepage:    "https://pulumi.io",
+		Repository:  "https://github.com/pulumi/pulumi-random",
+		Version:     "4.8.2",
+		P:           sdkv2.NewProvider(sdkv2randomprovider.New()),
 		Resources: map[string]*tfbridge.ResourceInfo{
 			"random_human_number": {Tok: randomResource(randomMod, "RandomHumanNumber")},
 		},
@@ -150,15 +157,13 @@ func RandomSDKProvider() tfbridge.ProviderInfo {
 	return info
 }
 
-func MuxedRandomProvider() []tfpf.Muxed {
+func MuxedRandomProvider() tfbridge.ProviderInfo {
 	sdk := RandomSDKProvider()
 	pf := RandomProvider()
-	for _, r := range pf.Resources {
+	for tf, r := range pf.Resources {
 		r.Tok = tokens.Type("muxedrandom:" + strings.TrimPrefix(string(r.Tok), "random:"))
+		sdk.Resources[tf] = r
 	}
-	pf.Name = sdk.Name
-	return []tfpf.Muxed{
-		{SDK: &sdk},
-		{PF: &pf},
-	}
+	sdk.P = tfpf.AugmentShimWithPF(context.Background(), sdk.P, pf.NewProvider())
+	return sdk
 }
diff --git a/pf/tests/provider_get_mapping_test.go b/pf/tests/provider_get_mapping_test.go
index a6538b56..19b9e510 100644
--- a/pf/tests/provider_get_mapping_test.go
+++ b/pf/tests/provider_get_mapping_test.go
@@ -25,9 +25,7 @@ import (
 
 	"github.com/pulumi/pulumi-terraform-bridge/pf/tests/internal/testprovider"
 	"github.com/pulumi/pulumi-terraform-bridge/pf/tfbridge"
-	"github.com/pulumi/pulumi-terraform-bridge/pf/tfgen"
 	tfbridge0 "github.com/pulumi/pulumi-terraform-bridge/v3/pkg/tfbridge"
-	tfgen0 "github.com/pulumi/pulumi-terraform-bridge/v3/pkg/tfgen"
 )
 
 func TestGetMapping(t *testing.T) {
@@ -64,17 +62,9 @@ func TestGetMapping(t *testing.T) {
 func TestMuxedGetMapping(t *testing.T) {
 	ctx := context.Background()
 
-	servers := testprovider.MuxedRandomProvider()
+	info := testprovider.MuxedRandomProvider()
 
-	_, _, schema, _, err := tfgen.UnstableMuxProviderInfo(ctx,
-		tfgen0.GeneratorOptions{}, "muxedrandom", servers...)
-	require.NoError(t, err)
-
-	schemaBytes, err := json.Marshal(schema)
-	require.NoError(t, err)
-
-	server, err := tfbridge.MakeMuxedServer(ctx, "random",
-		tfbridge.ProviderMetadata{PackageSchema: schemaBytes}, servers...)(nil)
+	server, err := tfbridge.MakeMuxedServer(ctx, genSDKSchema(t, info), info)(nil)
 	require.NoError(t, err)
 
 	req := func(key string) (context.Context, *pulumirpc.GetMappingRequest) {
diff --git a/pf/tests/schemas.go b/pf/tests/schemas.go
index a8e3d5bc..a40051f2 100644
--- a/pf/tests/schemas.go
+++ b/pf/tests/schemas.go
@@ -17,6 +17,7 @@ package tfbridgetests
 import (
 	"bytes"
 	"context"
+	"encoding/json"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -26,6 +27,8 @@ import (
 
 	tfpf "github.com/pulumi/pulumi-terraform-bridge/pf/tfbridge"
 	"github.com/pulumi/pulumi-terraform-bridge/pf/tfgen"
+	tfbridge0 "github.com/pulumi/pulumi-terraform-bridge/v3/pkg/tfbridge"
+	tfgen0 "github.com/pulumi/pulumi-terraform-bridge/v3/pkg/tfgen"
 )
 
 func genMetadata(t *testing.T, info tfpf.ProviderInfo) tfpf.ProviderMetadata {
@@ -51,3 +54,11 @@ func testSink(t *testing.T) diag.Sink {
 
 	return testSink
 }
+
+func genSDKSchema(t *testing.T, info tfbridge0.ProviderInfo) []byte {
+	pkg, err := tfgen0.GenerateSchema(info, testSink(t))
+	require.NoError(t, err)
+	bytes, err := json.MarshalIndent(pkg, "", "  ")
+	require.NoError(t, err)
+	return bytes
+}