refactor(provider): make googleFactory's gemini vs vertex dispatch testable

googleFactory contains the only non-trivial dispatch in the factory registry: it routes Gemini calls to gemini.NewClient and Vertex Model Garden calls (publisher != google) to vertexai.NewClient. Both arms reach out to real Google APIs, so neither was reachable from a unit test.

Lift the two inner constructors into package-level providerFactory variables (geminiClientFactory / vertexClientFactory). Tests can swap them via t.Cleanup and assert that vertexai.IsModelGardenConfig is consulted correctly:

  - plain google config -> gemini,

  - publisher=google still -> gemini (documented edge case),

  - publisher=anthropic -> vertex (Model Garden),

  - errors from either inner factory propagate unchanged.

Behaviour is unchanged. googleFactory goes from 0% to 100% coverage; package coverage 96.5% -> 97.1%.

Assisted-By: docker-agent

Author: dgageot

pkg/model/provider/google_factory_test.go

@@ -0,0 +1,134 @@
+package provider
+
+import (
+	"context"
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/docker/docker-agent/pkg/config/latest"
+	"github.com/docker/docker-agent/pkg/environment"
+	"github.com/docker/docker-agent/pkg/model/provider/options"
+)
+
+// withGoogleFactories swaps the gemini and vertex inner factories for the
+// duration of a test, restoring the production functions via t.Cleanup.
+func withGoogleFactories(t *testing.T, gemini, vertex providerFactory) {
+	t.Helper()
+	originalGemini, originalVertex := geminiClientFactory, vertexClientFactory
+	geminiClientFactory = gemini
+	vertexClientFactory = vertex
+	t.Cleanup(func() {
+		geminiClientFactory = originalGemini
+		vertexClientFactory = originalVertex
+	})
+}
+
+// TestGoogleFactory_RoutesGeminiByDefault verifies that a plain "google" model
+// (no Vertex Model Garden hints) is dispatched to the Gemini client.
+func TestGoogleFactory_RoutesGeminiByDefault(t *testing.T) {
+	withGoogleFactories(t,
+		tagFactory("gemini"),
+		func(_ context.Context, _ *latest.ModelConfig, _ environment.Provider, _ ...options.Opt) (Provider, error) {
+			t.Errorf("vertex factory should not be called for plain Gemini config")
+			return nil, errors.New("unreachable")
+		},
+	)
+
+	cfg := &latest.ModelConfig{Provider: "google", Model: "gemini-2.5-flash"}
+
+	p, err := googleFactory(t.Context(), cfg, environment.NewNoEnvProvider())
+	require.NoError(t, err)
+	fp, ok := p.(*fakeProvider)
+	require.True(t, ok)
+	assert.Equal(t, "gemini", fp.id)
+}
+
+// TestGoogleFactory_RoutesGeminiWhenPublisherIsGoogle covers the documented
+// edge case in vertexai.IsModelGardenConfig: publisher=google still routes to
+// Gemini (it's only a Model Garden config when publisher is non-google).
+func TestGoogleFactory_RoutesGeminiWhenPublisherIsGoogle(t *testing.T) {
+	withGoogleFactories(t,
+		tagFactory("gemini"),
+		func(_ context.Context, _ *latest.ModelConfig, _ environment.Provider, _ ...options.Opt) (Provider, error) {
+			t.Errorf("vertex factory must not be called when publisher=google")
+			return nil, errors.New("unreachable")
+		},
+	)
+
+	cfg := &latest.ModelConfig{
+		Provider:     "google",
+		Model:        "gemini-2.5-flash",
+		ProviderOpts: map[string]any{"publisher": "google"},
+	}
+
+	p, err := googleFactory(t.Context(), cfg, environment.NewNoEnvProvider())
+	require.NoError(t, err)
+	fp, ok := p.(*fakeProvider)
+	require.True(t, ok)
+	assert.Equal(t, "gemini", fp.id)
+}
+
+// TestGoogleFactory_RoutesVertexForModelGarden verifies that any non-Google
+// publisher routes through the Vertex Model Garden factory.
+func TestGoogleFactory_RoutesVertexForModelGarden(t *testing.T) {
+	withGoogleFactories(t,
+		func(_ context.Context, _ *latest.ModelConfig, _ environment.Provider, _ ...options.Opt) (Provider, error) {
+			t.Errorf("gemini factory must not be called for Model Garden config")
+			return nil, errors.New("unreachable")
+		},
+		tagFactory("vertex"),
+	)
+
+	cfg := &latest.ModelConfig{
+		Provider:     "google",
+		Model:        "claude-3-5-sonnet@20240620",
+		ProviderOpts: map[string]any{"publisher": "anthropic"},
+	}
+
+	p, err := googleFactory(t.Context(), cfg, environment.NewNoEnvProvider())
+	require.NoError(t, err)
+	fp, ok := p.(*fakeProvider)
+	require.True(t, ok)
+	assert.Equal(t, "vertex", fp.id)
+}
+
+// TestGoogleFactory_PropagatesGeminiError verifies that errors from the inner
+// gemini factory are surfaced unchanged.
+func TestGoogleFactory_PropagatesGeminiError(t *testing.T) {
+	sentinel := errors.New("gemini-fail")
+	withGoogleFactories(t,
+		func(_ context.Context, _ *latest.ModelConfig, _ environment.Provider, _ ...options.Opt) (Provider, error) {
+			return nil, sentinel
+		},
+		tagFactory("vertex"),
+	)
+
+	cfg := &latest.ModelConfig{Provider: "google", Model: "gemini-2.5-flash"}
+
+	_, err := googleFactory(t.Context(), cfg, environment.NewNoEnvProvider())
+	require.ErrorIs(t, err, sentinel)
+}
+
+// TestGoogleFactory_PropagatesVertexError verifies that errors from the inner
+// vertex factory are surfaced unchanged.
+func TestGoogleFactory_PropagatesVertexError(t *testing.T) {
+	sentinel := errors.New("vertex-fail")
+	withGoogleFactories(t,
+		tagFactory("gemini"),
+		func(_ context.Context, _ *latest.ModelConfig, _ environment.Provider, _ ...options.Opt) (Provider, error) {
+			return nil, sentinel
+		},
+	)
+
+	cfg := &latest.ModelConfig{
+		Provider:     "google",
+		Model:        "claude-3-5-sonnet@20240620",
+		ProviderOpts: map[string]any{"publisher": "anthropic"},
+	}
+
+	_, err := googleFactory(t.Context(), cfg, environment.NewNoEnvProvider())
+	require.ErrorIs(t, err, sentinel)
+}

pkg/model/provider/provider.go

@@ -306,11 +306,25 @@ func googleFactory(ctx context.Context, cfg *latest.ModelConfig, env environment
 	// Route non-Gemini models on Vertex AI (Model Garden) through the
 	// vertexai package, which picks the right endpoint per publisher.
 	if vertexai.IsModelGardenConfig(cfg) {
-		return vertexai.NewClient(ctx, cfg, env, opts...)
+		return vertexClientFactory(ctx, cfg, env, opts...)
 	}
-	return gemini.NewClient(ctx, cfg, env, opts...)
+	return geminiClientFactory(ctx, cfg, env, opts...)
 }
 
+// geminiClientFactory and vertexClientFactory are the inner constructors used
+// by googleFactory. They are package-level variables (rather than direct
+// references to gemini.NewClient / vertexai.NewClient) so that tests can swap
+// them with fakes via t.Cleanup and assert that googleFactory routes correctly
+// based on vertexai.IsModelGardenConfig — without spinning up real clients.
+var (
+	geminiClientFactory providerFactory = func(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+		return gemini.NewClient(ctx, cfg, env, opts...)
+	}
+	vertexClientFactory providerFactory = func(ctx context.Context, cfg *latest.ModelConfig, env environment.Provider, opts ...options.Opt) (Provider, error) {
+		return vertexai.NewClient(ctx, cfg, env, opts...)
+	}
+)
+
 func dmrFactory(ctx context.Context, cfg *latest.ModelConfig, _ environment.Provider, opts ...options.Opt) (Provider, error) {
 	return dmr.NewClient(ctx, cfg, opts...)
 }