implement mock iterators, improve tests

pkg/datastore/client.go

@@ -49,36 +49,34 @@ var (
 	}
 )
 
-// Config holds datastore client configuration.
-type Config struct {
-	// AuthConfig is passed to the auth package for authentication.
-	// Can be nil to use defaults.
-	AuthConfig *auth.Config
-
-	// APIURL is the base URL for the Datastore API.
-	// Defaults to production if empty.
-	APIURL string
-}
+// ClientOption is an option for a Datastore client.
+type ClientOption func(*clientOptionsInternal)
 
-// configKey is the key for storing Config in context.
-type configKey struct{}
+// clientOptionsInternal holds internal client configuration that can be modified by ClientOption.
+type clientOptionsInternal struct {
+	authConfig *auth.Config
+	logger     *slog.Logger
+	baseURL    string
+}
 
-// WithConfig returns a new context with the given datastore config.
-// This also sets the auth config if provided.
-func WithConfig(ctx context.Context, cfg *Config) context.Context {
-	if cfg.AuthConfig != nil {
-		ctx = auth.WithConfig(ctx, cfg.AuthConfig)
+// WithEndpoint returns a ClientOption that sets the API base URL.
+func WithEndpoint(url string) ClientOption {
+	return func(o *clientOptionsInternal) {
+		o.baseURL = url
 	}
-	return context.WithValue(ctx, configKey{}, cfg)
 }
 
-// getConfig retrieves the datastore config from context, or returns defaults.
-func getConfig(ctx context.Context) *Config {
-	if cfg, ok := ctx.Value(configKey{}).(*Config); ok && cfg != nil {
-		return cfg
+// WithLogger returns a ClientOption that sets the logger.
+func WithLogger(logger *slog.Logger) ClientOption {
+	return func(o *clientOptionsInternal) {
+		o.logger = logger
 	}
-	return &Config{
-		APIURL: defaultAPIURL,
+}
+
+// WithAuth returns a ClientOption that sets the authentication configuration.
+func WithAuth(cfg *auth.Config) ClientOption {
+	return func(o *clientOptionsInternal) {
+		o.authConfig = cfg
 	}
 }
 
@@ -93,37 +91,52 @@ type Client struct {
 
 // NewClient creates a new Datastore client.
 // If projectID is empty, it will be fetched from the GCP metadata server.
-// Configuration can be provided via WithConfig in the context.
-func NewClient(ctx context.Context, projectID string) (*Client, error) {
-	return NewClientWithDatabase(ctx, projectID, "")
+// Options can be provided to configure the client.
+func NewClient(ctx context.Context, projectID string, opts ...ClientOption) (*Client, error) {
+	return NewClientWithDatabase(ctx, projectID, "", opts...)
 }
 
 // NewClientWithDatabase creates a new Datastore client with a specific database.
-// Configuration can be provided via WithConfig in the context.
-func NewClientWithDatabase(ctx context.Context, projID, dbID string) (*Client, error) {
-	logger := slog.Default()
-	cfg := getConfig(ctx)
+// Options can be provided to configure the client.
+func NewClientWithDatabase(ctx context.Context, projID, dbID string, opts ...ClientOption) (*Client, error) {
+	// Apply default internal options
+	options := &clientOptionsInternal{
+		baseURL: defaultAPIURL,
+		logger:  slog.Default(),
+	}
+
+	// Apply provided options
+	for _, opt := range opts {
+		opt(options)
+	}
 
+	// --- Existing NewClientWithDatabase logic starts here ---
 	if projID == "" {
+		// Inject auth config into context before fetching project ID
+		fetchCtx := ctx
+		if options.authConfig != nil {
+			fetchCtx = auth.WithConfig(ctx, options.authConfig)
+		}
+
 		if !testing.Testing() {
-			logger.InfoContext(ctx, "project ID not provided, fetching from metadata server")
+			options.logger.InfoContext(ctx, "project ID not provided, fetching from metadata server")
 		}
-		pid, err := auth.ProjectID(ctx)
+		pid, err := auth.ProjectID(fetchCtx)
 		if err != nil {
-			logger.ErrorContext(ctx, "failed to get project ID from metadata server", "error", err)
+			options.logger.ErrorContext(ctx, "failed to get project ID from metadata server", "error", err)
 			return nil, fmt.Errorf("project ID required: %w", err)
 		}
 		projID = pid
 		if !testing.Testing() {
-			logger.InfoContext(ctx, "fetched project ID from metadata server", "project_id", projID)
+			options.logger.InfoContext(ctx, "fetched project ID from metadata server", "project_id", projID)
 		}
 	}
 
 	if !testing.Testing() {
-		logger.InfoContext(ctx, "creating datastore client", "project_id", projID, "database_id", dbID)
+		options.logger.InfoContext(ctx, "creating datastore client", "project_id", projID, "database_id", dbID)
 	}
 
-	baseURL := cfg.APIURL
+	baseURL := options.baseURL
 	if baseURL == "" {
 		baseURL = defaultAPIURL
 	}
@@ -132,8 +145,8 @@ func NewClientWithDatabase(ctx context.Context, projID, dbID string) (*Client, e
 		projectID:  projID,
 		databaseID: dbID,
 		baseURL:    baseURL,
-		authConfig: cfg.AuthConfig,
-		logger:     logger,
+		authConfig: options.authConfig, // Use authConfig from options
+		logger:     options.logger,     // Use logger from options
 	}, nil
 }
 

pkg/datastore/client_test.go

@@ -59,16 +59,18 @@ func TestNewClientWithDatabase(t *testing.T) {
 	}))
 	defer apiServer.Close()
 
-	ctx := datastore.WithConfig(context.Background(), &datastore.Config{
-		APIURL: apiServer.URL,
-		AuthConfig: &auth.Config{
-			MetadataURL: metadataServer.URL,
-			SkipADC:     true,
-		},
-	})
-
+	authConfig := &auth.Config{
+		MetadataURL: metadataServer.URL,
+		SkipADC:     true,
+	}
 	// Test with explicit databaseID
-	client, err := datastore.NewClientWithDatabase(ctx, "test-project", "custom-db")
+	client, err := datastore.NewClientWithDatabase(
+		context.Background(),
+		"test-project",
+		"custom-db",
+		datastore.WithEndpoint(apiServer.URL),
+		datastore.WithAuth(authConfig),
+	)
 	if err != nil {
 		t.Fatalf("NewClientWithDatabase failed: %v", err)
 	}
@@ -77,23 +79,6 @@ func TestNewClientWithDatabase(t *testing.T) {
 	}
 }
 
-func TestWithConfig(t *testing.T) {
-	// Test that config can be set in context
-	cfg := &datastore.Config{
-		APIURL: "http://test",
-		AuthConfig: &auth.Config{
-			MetadataURL: "http://metadata",
-			SkipADC:     true,
-		},
-	}
-	ctx := datastore.WithConfig(context.Background(), cfg)
-
-	// Context should be non-nil
-	if ctx == nil {
-		t.Fatal("expected non-nil context")
-	}
-}
-
 func TestNewClientWithDatabaseEmptyProjectID(t *testing.T) {
 	// Setup mock servers
 	metadataServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -131,16 +116,18 @@ func TestNewClientWithDatabaseEmptyProjectID(t *testing.T) {
 	}))
 	defer apiServer.Close()
 
-	ctx := datastore.WithConfig(context.Background(), &datastore.Config{
-		APIURL: apiServer.URL,
-		AuthConfig: &auth.Config{
-			MetadataURL: metadataServer.URL,
-			SkipADC:     true,
-		},
-	})
-
+	authConfig := &auth.Config{
+		MetadataURL: metadataServer.URL,
+		SkipADC:     true,
+	}
 	// Test with empty projectID - should fetch from metadata
-	client, err := datastore.NewClientWithDatabase(ctx, "", "my-db")
+	client, err := datastore.NewClientWithDatabase(
+		context.Background(),
+		"",
+		"my-db",
+		datastore.WithEndpoint(apiServer.URL),
+		datastore.WithAuth(authConfig),
+	)
 	if err != nil {
 		t.Fatalf("NewClientWithDatabase with empty projectID failed: %v", err)
 	}
@@ -184,16 +171,18 @@ func TestNewClientWithDatabaseProjectIDFetchFailure(t *testing.T) {
 	}))
 	defer apiServer.Close()
 
-	ctx := datastore.WithConfig(context.Background(), &datastore.Config{
-		APIURL: apiServer.URL,
-		AuthConfig: &auth.Config{
-			MetadataURL: metadataServer.URL,
-			SkipADC:     true,
-		},
-	})
-
+	authConfig := &auth.Config{
+		MetadataURL: metadataServer.URL,
+		SkipADC:     true,
+	}
 	// Test with empty projectID and failing metadata server
-	client, err := datastore.NewClientWithDatabase(ctx, "", "my-db")
+	client, err := datastore.NewClientWithDatabase(
+		context.Background(),
+		"",
+		"my-db",
+		datastore.WithEndpoint(apiServer.URL),
+		datastore.WithAuth(authConfig),
+	)
 	if err == nil {
 		t.Fatal("expected error when projectID fetch fails, got nil")
 	}

pkg/datastore/entity_coverage_test.go

@@ -314,10 +314,10 @@ func TestNewClientWithDatabase_Coverage(t *testing.T) {
 	}))
 	defer apiServer.Close()
 
-	testCtx := TestConfig(ctx, metadataServer.URL, apiServer.URL)
+	opts := TestConfig(ctx, metadataServer.URL, apiServer.URL)
 
 	// Test with empty project (error case)
-	_, err := NewClientWithDatabase(testCtx, "", "test-db")
+	_, err := NewClientWithDatabase(ctx, "", "test-db", opts...)
 	if err == nil {
 		t.Error("Expected error for empty project ID, got nil")
 	}