Inject mock more directly

provider/internal/image.go

@@ -61,7 +61,9 @@ var _imageExamples string
 var _migration string
 
 // Image is a Docker image build using buildkit.
-type Image struct{}
+type Image struct {
+	docker Client
+}
 
 // Annotate provides a description of the Image resource.
 func (i *Image) Annotate(a infer.Annotator) {
@@ -330,12 +332,13 @@ func (is *ImageState) Annotate(a infer.Annotator) {
 // client produces a CLI client scoped to this resource and layered on top of
 // any host-level credentials.
 func (i *Image) client(ctx context.Context, state ImageState, args ImageArgs) (Client, error) {
-	cfg := infer.GetConfig[Config](ctx)
-
-	if cli, ok := ctx.Value(_mockClientKey).(Client); ok {
-		return cli, nil
+	// Use our mock client, if it's set.
+	if i.docker != nil {
+		return i.docker, nil
 	}
 
+	cfg := infer.GetConfig[Config](ctx)
+
 	// We prefer auth from args, the provider, and state in that order. We
 	// build a slice in reverse order because wrap() will overwrite earlier
 	// entries with later ones.

provider/internal/index.go

@@ -46,7 +46,9 @@ var (
 var _indexExamples string
 
 // Index is an OCI index or manifest list on a remote registry.
-type Index struct{}
+type Index struct {
+	docker Client
+}
 
 // IndexArgs instantiate an Index.
 type IndexArgs struct {
@@ -351,12 +353,13 @@ func (i *Index) client(
 	_ IndexState,
 	args IndexArgs,
 ) (Client, error) {
-	cfg := infer.GetConfig[Config](ctx)
-
-	if cli, ok := ctx.Value(_mockClientKey).(Client); ok {
-		return cli, nil
+	// Use our mock client, if it's set.
+	if i.docker != nil {
+		return i.docker, nil
 	}
 
+	cfg := infer.GetConfig[Config](ctx)
+
 	// We prefer auth from args, the provider, and state in that order. We
 	// build a slice in reverse order because wrap() will overwrite earlier
 	// entries with later ones.

provider/internal/provider.go

@@ -46,9 +46,6 @@ type Config struct {
 	host *host
 }
 
-// _mockClientKey is used by tests to inject a mock Docker client.
-var _mockClientKey any = "mock-client"
-
 // Annotate provides user-facing descriptions and defaults for Config's fields.
 func (c *Config) Annotate(a infer.Annotator) {
 	a.Describe(&c.Host, "The build daemon's address.")
@@ -66,7 +63,7 @@ func (c *Config) Configure(ctx context.Context) error {
 }
 
 // NewBuildxProvider returns a new buildx provider.
-func NewBuildxProvider() provider.Provider {
+func NewBuildxProvider(mock Client) provider.Provider {
 	prov := infer.Provider(
 		infer.Options{
 			Metadata: pschema.Metadata{
@@ -114,13 +111,13 @@ func NewBuildxProvider() provider.Provider {
 				},
 			},
 			Resources: []infer.InferredResource{
-				infer.Resource[*Image](),
-				infer.Resource[*Index](),
+				infer.Resource(&Image{docker: mock}),
+				infer.Resource(&Index{docker: mock}),
 			},
 			ModuleMap: map[tokens.ModuleName]tokens.ModuleName{
 				"internal": "index",
 			},
-			Config: infer.Config[*Config](),
+			Config: infer.Config(&Config{}),
 		},
 	)
 
@@ -152,7 +149,7 @@ func diffConfigIgnoreInternal(
 
 // Schema returns our package specification.
 func Schema(ctx context.Context, version string) schema.PackageSpec {
-	p := NewBuildxProvider()
+	p := NewBuildxProvider(nil)
 	spec, err := provider.GetSchema(ctx, "docker-build", version, p)
 	contract.AssertNoErrorf(err, "missing schema")
 	return spec

provider/internal/provider_test.go

@@ -25,7 +25,6 @@ import (
 	provider "github.com/pulumi/pulumi-go-provider"
 	"github.com/pulumi/pulumi-go-provider/infer"
 	"github.com/pulumi/pulumi-go-provider/integration"
-	mwcontext "github.com/pulumi/pulumi-go-provider/middleware/context"
 	"github.com/pulumi/pulumi/sdk/v3/go/common/tokens"
 )
 
@@ -79,14 +78,7 @@ func (annotator) SetResourceDeprecationMessage(_ string)      {}
 func newServer(ctx context.Context, t *testing.T, client Client) integration.Server {
 	t.Helper()
 
-	p := NewBuildxProvider()
-
-	// Inject a mock client if provided.
-	if client != nil {
-		p = mwcontext.Wrap(p, func(ctx context.Context) context.Context {
-			return context.WithValue(ctx, _mockClientKey, client)
-		})
-	}
+	p := NewBuildxProvider(client)
 
 	s, err := integration.NewServer(
 		ctx,