From f2e43df5af95fadabab04aa25f287af0c8e6fb58 Mon Sep 17 00:00:00 2001
From: Bryce Lampe <bryce@pulumi.com>
Date: Fri, 26 Apr 2024 11:49:50 -0700
Subject: [PATCH] fix tests

---
 provider/internal/client_test.go   | 19 +++++++++++++++----
 provider/internal/provider_test.go |  7 +++++--
 provider/provider_test.go          | 28 +++++++++++++++++++++++-----
 3 files changed, 43 insertions(+), 11 deletions(-)

diff --git a/provider/internal/client_test.go b/provider/internal/client_test.go
index 8e7f00f..a1cdb3f 100644
--- a/provider/internal/client_test.go
+++ b/provider/internal/client_test.go
@@ -17,6 +17,8 @@ package internal
 import (
 	"bytes"
 	"context"
+	"io"
+	"log/slog"
 	"os"
 	"path/filepath"
 	"testing"
@@ -352,13 +354,23 @@ func TestNormalizeReference(t *testing.T) {
 	}
 }
 
+//nolint:paralleltest // Overrides default logger.
 func TestBuildError(t *testing.T) {
-	t.Parallel()
-
 	if os.Getenv("CI") != "" {
 		t.Skip("flaky on CI for some reason")
 	}
 
+	l := slog.Default()
+	defer slog.SetDefault(l)
+
+	// Override go-provider's default logger to capture and tee to stdout.
+	logger := &bytes.Buffer{}
+	slog.SetDefault(
+		slog.New(
+			slog.NewTextHandler(io.MultiWriter(logger, os.Stdout), nil),
+		),
+	)
+
 	exampleContext := &BuildContext{Context: Context{Location: "../../examples/app"}}
 
 	args := ImageArgs{
@@ -367,7 +379,6 @@ func TestBuildError(t *testing.T) {
 			Inline: "FROM alpine\nRUN echo hello\nRUN badcmd",
 		},
 	}
-	logged := bytes.Buffer{}
 
 	ctx := context.Background()
 	cli := testcli(t, true)
@@ -384,7 +395,7 @@ func TestBuildError(t *testing.T) {
 	}
 
 	for _, want := range want {
-		assert.Contains(t, logged.String(), want)
+		assert.Contains(t, logger.String(), want)
 	}
 	assert.ErrorContains(t, err,
 		`process "/bin/sh -c badcmd" did not complete successfully: exit code: 127`,
diff --git a/provider/internal/provider_test.go b/provider/internal/provider_test.go
index fef1429..480897f 100644
--- a/provider/internal/provider_test.go
+++ b/provider/internal/provider_test.go
@@ -59,7 +59,10 @@ func TestAnnotate(t *testing.T) {
 func TestSchema(t *testing.T) {
 	t.Parallel()
 
-	Schema(context.Background(), "v4")
+	s := newServer(nil)
+
+	_, err := s.GetSchema(provider.GetSchemaRequest{Version: 0})
+	assert.NoError(t, err)
 }
 
 type annotator struct{}
@@ -78,5 +81,5 @@ func newServer(client Client) integration.Server {
 		})
 	}
 
-	return integration.NewServer("docker", semver.Version{Major: 4}, p)
+	return integration.NewServer("docker-build", semver.Version{Major: 0}, p)
 }
diff --git a/provider/provider_test.go b/provider/provider_test.go
index 0e18dcb..8ade083 100644
--- a/provider/provider_test.go
+++ b/provider/provider_test.go
@@ -24,7 +24,10 @@ import (
 	emptypb "google.golang.org/protobuf/types/known/emptypb"
 	"google.golang.org/protobuf/types/known/structpb"
 
-	pulumirpc "github.com/pulumi/pulumi/sdk/v3/proto/go"
+	provider "github.com/pulumi/pulumi-go-provider"
+	"github.com/pulumi/pulumi-go-provider/integration"
+	"github.com/pulumi/pulumi/sdk/v3/go/common/resource/plugin"
+	rpc "github.com/pulumi/pulumi/sdk/v3/proto/go"
 )
 
 // TestConfigure checks backwards-compatibility with SDKs that still send
@@ -34,8 +37,6 @@ import (
 func TestConfigure(t *testing.T) {
 	t.Parallel()
 
-	ctx := context.Background()
-
 	p, err := New(nil)
 	require.NoError(t, err)
 
@@ -44,8 +45,25 @@ func TestConfigure(t *testing.T) {
 	})
 	require.NoError(t, err)
 
-	_, err = p.Configure(ctx, &pulumirpc.ConfigureRequest{
-		Args: args,
+	// Ideally we would just call p.Configure directly, but we need the
+	// integration server to inject runtime info for us.
+	argsMap, err := plugin.UnmarshalProperties(args, plugin.MarshalOptions{})
+	require.NoError(t, err)
+	s := integration.NewServer("docker-build", semver.Version{Major: 0}, provider.Provider{
+		// Roundabout way to get the integration server to invoke our outermost
+		// Configure RPC endpoint.
+		Configure: func(ctx context.Context, req provider.ConfigureRequest) error {
+			args, err := plugin.MarshalProperties(req.Args, plugin.MarshalOptions{})
+			require.NoError(t, err)
+			_, err = p.Configure(ctx, &rpc.ConfigureRequest{
+				Variables: req.Variables,
+				Args:      args,
+			})
+			return err
+		},
+	})
+	err = s.Configure(provider.ConfigureRequest{
+		Args: argsMap,
 	})
 
 	assert.NoError(t, err)