Return primary provider server from 'provider.ProtoV5ProviderServerFa…

…ctory'.

internal/acctest/acctest.go

@@ -110,7 +110,7 @@ func protoV5ProviderFactoriesInit(providerNames ...string) map[string]func() (tf
 
 	for _, name := range providerNames {
 		factories[name] = func() (tfprotov5.ProviderServer, error) {
-			providerServerFactory, err := provider.ProtoV5ProviderServerFactory(context.Background())
+			providerServerFactory, _, err := provider.ProtoV5ProviderServerFactory(context.Background())
 
 			if err != nil {
 				return nil, err

internal/provider/factory.go

@@ -6,16 +6,18 @@ import (
 	"github.com/hashicorp/terraform-plugin-framework/providerserver"
 	"github.com/hashicorp/terraform-plugin-go/tfprotov5"
 	"github.com/hashicorp/terraform-plugin-mux/tf5muxserver"
+	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
 	"github.com/hashicorp/terraform-provider-aws/internal/provider/fwprovider"
 )
 
 // ProtoV5ProviderServerFactory returns a muxed terraform-plugin-go protocol v5 provider factory function.
 // This factory function is suitable for use with the terraform-plugin-go Serve function.
-func ProtoV5ProviderServerFactory(ctx context.Context) (func() tfprotov5.ProviderServer, error) {
+// The primary (Plugin SDK) provider server is also returned (useful for testing).
+func ProtoV5ProviderServerFactory(ctx context.Context) (func() tfprotov5.ProviderServer, *schema.Provider, error) {
 	primary, err := New(ctx)
 
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	servers := []func() tfprotov5.ProviderServer{
@@ -26,8 +28,8 @@ func ProtoV5ProviderServerFactory(ctx context.Context) (func() tfprotov5.Provide
 	muxServer, err := tf5muxserver.NewMuxServer(ctx, servers...)
 
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
-	return muxServer.ProviderServer, nil
+	return muxServer.ProviderServer, primary, nil
 }

main.go

@@ -13,7 +13,7 @@ func main() {
 	debugFlag := flag.Bool("debug", false, "Start provider in debug mode.")
 	flag.Parse()
 
-	serverFactory, err := provider.ProtoV5ProviderServerFactory(context.Background())
+	serverFactory, _, err := provider.ProtoV5ProviderServerFactory(context.Background())
 
 	if err != nil {
 		log.Fatal(err)