Skip to content

Commit 173defd

Browse files
tstrombergtstromberg
authored
Merge pull request #11 from tstromberg/main
implement mock iterators, improve tests
2 parents f92b1f3 + 8786fe6 commit 173defd

14 files changed

+1061
-685
lines changed

‎pkg/datastore/client.go‎

Lines changed: 52 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -49,36 +49,34 @@ var (
4949
}
5050
)
5151

52-
// Config holds datastore client configuration.
53-
type Config struct {
54-
// AuthConfig is passed to the auth package for authentication.
55-
// Can be nil to use defaults.
56-
AuthConfig *auth.Config
57-
58-
// APIURL is the base URL for the Datastore API.
59-
// Defaults to production if empty.
60-
APIURL string
61-
}
52+
// ClientOption is an option for a Datastore client.
53+
type ClientOption func(*clientOptionsInternal)
6254

63-
// configKey is the key for storing Config in context.
64-
type configKey struct{}
55+
// clientOptionsInternal holds internal client configuration that can be modified by ClientOption.
56+
type clientOptionsInternal struct {
57+
authConfig *auth.Config
58+
logger *slog.Logger
59+
baseURL string
60+
}
6561

66-
// WithConfig returns a new context with the given datastore config.
67-
// This also sets the auth config if provided.
68-
func WithConfig(ctx context.Context, cfg *Config) context.Context {
69-
if cfg.AuthConfig != nil {
70-
ctx = auth.WithConfig(ctx, cfg.AuthConfig)
62+
// WithEndpoint returns a ClientOption that sets the API base URL.
63+
func WithEndpoint(url string) ClientOption {
64+
return func(o *clientOptionsInternal) {
65+
o.baseURL = url
7166
}
72-
return context.WithValue(ctx, configKey{}, cfg)
7367
}
7468

75-
// getConfig retrieves the datastore config from context, or returns defaults.
76-
func getConfig(ctx context.Context) *Config {
77-
if cfg, ok := ctx.Value(configKey{}).(*Config); ok && cfg != nil {
78-
return cfg
69+
// WithLogger returns a ClientOption that sets the logger.
70+
func WithLogger(logger *slog.Logger) ClientOption {
71+
return func(o *clientOptionsInternal) {
72+
o.logger = logger
7973
}
80-
return &Config{
81-
APIURL: defaultAPIURL,
74+
}
75+
76+
// WithAuth returns a ClientOption that sets the authentication configuration.
77+
func WithAuth(cfg *auth.Config) ClientOption {
78+
return func(o *clientOptionsInternal) {
79+
o.authConfig = cfg
8280
}
8381
}
8482

@@ -93,37 +91,52 @@ type Client struct {
9391

9492
// NewClient creates a new Datastore client.
9593
// If projectID is empty, it will be fetched from the GCP metadata server.
96-
// Configuration can be provided via WithConfig in the context.
97-
func NewClient(ctx context.Context, projectID string) (*Client, error) {
98-
return NewClientWithDatabase(ctx, projectID, "")
94+
// Options can be provided to configure the client.
95+
func NewClient(ctx context.Context, projectID string, opts ...ClientOption) (*Client, error) {
96+
return NewClientWithDatabase(ctx, projectID, "", opts...)
9997
}
10098

10199
// NewClientWithDatabase creates a new Datastore client with a specific database.
102-
// Configuration can be provided via WithConfig in the context.
103-
func NewClientWithDatabase(ctx context.Context, projID, dbID string) (*Client, error) {
104-
logger := slog.Default()
105-
cfg := getConfig(ctx)
100+
// Options can be provided to configure the client.
101+
func NewClientWithDatabase(ctx context.Context, projID, dbID string, opts ...ClientOption) (*Client, error) {
102+
// Apply default internal options
103+
options := &clientOptionsInternal{
104+
baseURL: defaultAPIURL,
105+
logger: slog.Default(),
106+
}
107+
108+
// Apply provided options
109+
for _, opt := range opts {
110+
opt(options)
111+
}
106112

113+
// --- Existing NewClientWithDatabase logic starts here ---
107114
if projID == "" {
115+
// Inject auth config into context before fetching project ID
116+
fetchCtx := ctx
117+
if options.authConfig != nil {
118+
fetchCtx = auth.WithConfig(ctx, options.authConfig)
119+
}
120+
108121
if !testing.Testing() {
109-
logger.InfoContext(ctx, "project ID not provided, fetching from metadata server")
122+
options.logger.InfoContext(ctx, "project ID not provided, fetching from metadata server")
110123
}
111-
pid, err := auth.ProjectID(ctx)
124+
pid, err := auth.ProjectID(fetchCtx)
112125
if err != nil {
113-
logger.ErrorContext(ctx, "failed to get project ID from metadata server", "error", err)
126+
options.logger.ErrorContext(ctx, "failed to get project ID from metadata server", "error", err)
114127
return nil, fmt.Errorf("project ID required: %w", err)
115128
}
116129
projID = pid
117130
if !testing.Testing() {
118-
logger.InfoContext(ctx, "fetched project ID from metadata server", "project_id", projID)
131+
options.logger.InfoContext(ctx, "fetched project ID from metadata server", "project_id", projID)
119132
}
120133
}
121134

122135
if !testing.Testing() {
123-
logger.InfoContext(ctx, "creating datastore client", "project_id", projID, "database_id", dbID)
136+
options.logger.InfoContext(ctx, "creating datastore client", "project_id", projID, "database_id", dbID)
124137
}
125138

126-
baseURL := cfg.APIURL
139+
baseURL := options.baseURL
127140
if baseURL == "" {
128141
baseURL = defaultAPIURL
129142
}
@@ -132,8 +145,8 @@ func NewClientWithDatabase(ctx context.Context, projID, dbID string) (*Client, e
132145
projectID: projID,
133146
databaseID: dbID,
134147
baseURL: baseURL,
135-
authConfig: cfg.AuthConfig,
136-
logger: logger,
148+
authConfig: options.authConfig, // Use authConfig from options
149+
logger: options.logger, // Use logger from options
137150
}, nil
138151
}
139152

‎pkg/datastore/client_test.go‎

Lines changed: 33 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -59,16 +59,18 @@ func TestNewClientWithDatabase(t *testing.T) {
5959
}))
6060
defer apiServer.Close()
6161

62-
ctx := datastore.WithConfig(context.Background(), &datastore.Config{
63-
APIURL: apiServer.URL,
64-
AuthConfig: &auth.Config{
65-
MetadataURL: metadataServer.URL,
66-
SkipADC: true,
67-
},
68-
})
69-
62+
authConfig := &auth.Config{
63+
MetadataURL: metadataServer.URL,
64+
SkipADC: true,
65+
}
7066
// Test with explicit databaseID
71-
client, err := datastore.NewClientWithDatabase(ctx, "test-project", "custom-db")
67+
client, err := datastore.NewClientWithDatabase(
68+
context.Background(),
69+
"test-project",
70+
"custom-db",
71+
datastore.WithEndpoint(apiServer.URL),
72+
datastore.WithAuth(authConfig),
73+
)
7274
if err != nil {
7375
t.Fatalf("NewClientWithDatabase failed: %v", err)
7476
}
@@ -77,23 +79,6 @@ func TestNewClientWithDatabase(t *testing.T) {
7779
}
7880
}
7981

80-
func TestWithConfig(t *testing.T) {
81-
// Test that config can be set in context
82-
cfg := &datastore.Config{
83-
APIURL: "http://test",
84-
AuthConfig: &auth.Config{
85-
MetadataURL: "http://metadata",
86-
SkipADC: true,
87-
},
88-
}
89-
ctx := datastore.WithConfig(context.Background(), cfg)
90-
91-
// Context should be non-nil
92-
if ctx == nil {
93-
t.Fatal("expected non-nil context")
94-
}
95-
}
96-
9782
func TestNewClientWithDatabaseEmptyProjectID(t *testing.T) {
9883
// Setup mock servers
9984
metadataServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
@@ -131,16 +116,18 @@ func TestNewClientWithDatabaseEmptyProjectID(t *testing.T) {
131116
}))
132117
defer apiServer.Close()
133118

134-
ctx := datastore.WithConfig(context.Background(), &datastore.Config{
135-
APIURL: apiServer.URL,
136-
AuthConfig: &auth.Config{
137-
MetadataURL: metadataServer.URL,
138-
SkipADC: true,
139-
},
140-
})
141-
119+
authConfig := &auth.Config{
120+
MetadataURL: metadataServer.URL,
121+
SkipADC: true,
122+
}
142123
// Test with empty projectID - should fetch from metadata
143-
client, err := datastore.NewClientWithDatabase(ctx, "", "my-db")
124+
client, err := datastore.NewClientWithDatabase(
125+
context.Background(),
126+
"",
127+
"my-db",
128+
datastore.WithEndpoint(apiServer.URL),
129+
datastore.WithAuth(authConfig),
130+
)
144131
if err != nil {
145132
t.Fatalf("NewClientWithDatabase with empty projectID failed: %v", err)
146133
}
@@ -184,16 +171,18 @@ func TestNewClientWithDatabaseProjectIDFetchFailure(t *testing.T) {
184171
}))
185172
defer apiServer.Close()
186173

187-
ctx := datastore.WithConfig(context.Background(), &datastore.Config{
188-
APIURL: apiServer.URL,
189-
AuthConfig: &auth.Config{
190-
MetadataURL: metadataServer.URL,
191-
SkipADC: true,
192-
},
193-
})
194-
174+
authConfig := &auth.Config{
175+
MetadataURL: metadataServer.URL,
176+
SkipADC: true,
177+
}
195178
// Test with empty projectID and failing metadata server
196-
client, err := datastore.NewClientWithDatabase(ctx, "", "my-db")
179+
client, err := datastore.NewClientWithDatabase(
180+
context.Background(),
181+
"",
182+
"my-db",
183+
datastore.WithEndpoint(apiServer.URL),
184+
datastore.WithAuth(authConfig),
185+
)
197186
if err == nil {
198187
t.Fatal("expected error when projectID fetch fails, got nil")
199188
}

‎pkg/datastore/entity_coverage_test.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -314,10 +314,10 @@ func TestNewClientWithDatabase_Coverage(t *testing.T) {
314314
}))
315315
defer apiServer.Close()
316316

317-
testCtx := TestConfig(ctx, metadataServer.URL, apiServer.URL)
317+
opts := TestConfig(ctx, metadataServer.URL, apiServer.URL)
318318

319319
// Test with empty project (error case)
320-
_, err := NewClientWithDatabase(testCtx, "", "test-db")
320+
_, err := NewClientWithDatabase(ctx, "", "test-db", opts...)
321321
if err == nil {
322322
t.Error("Expected error for empty project ID, got nil")
323323
}

0 commit comments

Comments
 (0)