Less abstraction for integration tests

Makefile

@@ -15,7 +15,7 @@ integration:
 	@echo ""
 	@echo "==> Checking if test database exists..."
 	@if gcloud firestore databases describe --database=$(DS9_TEST_DATABASE) --project=$(DS9_TEST_PROJECT) >/dev/null 2>&1; then \
-		echo "    Database already exists, skipping creation"; \
+		echo "    Database already exists, using existing database"; \
 	else \
 		echo "    Database does not exist, creating..."; \
 		if ! gcloud firestore databases create --database=$(DS9_TEST_DATABASE) \
@@ -31,17 +31,21 @@ integration:
 			exit 1; \
 		fi; \
 		echo "    Database created successfully"; \
+		echo "    Waiting 10 seconds for database to propagate..."; \
+		sleep 10; \
 	fi
 	@echo ""
 	@echo "==> Running integration tests..."
-	@DS9_TEST_PROJECT=$(DS9_TEST_PROJECT) go test -v -race -tags=integration -timeout=5m ./... || \
-		(echo ""; echo "==> Tests failed, cleaning up..."; \
-		 gcloud firestore databases delete --database=$(DS9_TEST_DATABASE) --project=$(DS9_TEST_PROJECT) --quiet 2>/dev/null; \
-		 exit 1)
+	@DS9_TEST_PROJECT=$(DS9_TEST_PROJECT) go test -v -race -timeout=5m ./...
 	@echo ""
-	@echo "==> Cleaning up test database..."
-	@gcloud firestore databases delete --database=$(DS9_TEST_DATABASE) --project=$(DS9_TEST_PROJECT) --quiet
 	@echo "==> Integration tests complete!"
+	@echo "Note: Database $(DS9_TEST_DATABASE) is retained for reuse"
+	@echo "      Run 'make clean-integration-db' to delete it"
+
+clean-integration-db:
+	@echo "==> Deleting integration test database..."
+	@gcloud firestore databases delete --database=$(DS9_TEST_DATABASE) --project=$(DS9_TEST_PROJECT) --quiet
+	@echo "==> Database deleted successfully"
 
 lint:
 	go vet ./...

TESTING.md

@@ -8,11 +8,11 @@ Run the standard test suite with mock Datastore:
 make test
 ```
 
-This uses `ds9mock` for in-memory testing without requiring GCP credentials.
+This runs all tests including integration tests against an in-memory mock server. No GCP credentials required.
 
 ## Integration Tests
 
-Integration tests run against real Google Cloud Datastore and automatically manage database lifecycle.
+Integration tests automatically use the mock server by default, but can run against real Google Cloud Datastore when `DS9_TEST_PROJECT` is set.
 
 ### Prerequisites
 
@@ -40,9 +40,10 @@ make integration
 ```
 
 This will automatically:
-1. Create a temporary Datastore database (`ds9-test`)
-2. Run the full integration test suite
-3. Delete the temporary database (even if tests fail)
+1. Check if the test database (`ds9-test`) exists, or create it if needed
+2. If creating a new database, wait 10 seconds for it to propagate (GCP needs time to make the database available)
+3. Run the full integration test suite (including cleanup of test entities)
+4. Retain the database for reuse in subsequent test runs
 
 **Customization:**
 ```bash
@@ -56,29 +57,39 @@ make integration DS9_TEST_DATABASE=my-test-db
 make integration DS9_TEST_LOCATION=europe-west1
 ```
 
+### How It Works
+
+- **Without `DS9_TEST_PROJECT`**: Tests run against an in-memory mock server (fast, no GCP needed)
+- **With `DS9_TEST_PROJECT`**: Tests run against real Google Cloud Datastore (requires GCP credentials)
+
+The same test code runs in both modes, ensuring the mock accurately represents real Datastore behavior.
+
 ### What Gets Tested
 
 - **Basic Operations**: Put, Get, Update, Delete
 - **Batch Operations**: PutMulti, GetMulti, DeleteMulti
 - **Transactions**: Read-modify-write operations
 - **Queries**: KeysOnly queries with limits
+- **Cleanup**: DeleteAllByKind operation
 
 ### Test Data
 
-Test entities use the kind `DS9IntegrationTest` with unique timestamp-based names to avoid conflicts. The test database (`ds9-test`) is automatically created before tests and deleted after, ensuring complete cleanup.
+Test entities use the kind `DS9IntegrationTest` with unique timestamp-based names to avoid conflicts. Each test run creates entities, and the final cleanup test deletes all entities of this kind.
+
+### Database Cleanup
 
-### Manual Cleanup
+The test database is retained between test runs for performance (database creation takes several minutes). Test entities are automatically cleaned up at the end of each test run.
 
-The Makefile automatically cleans up the test database, even if tests fail. If you need to manually clean up:
+To manually delete the test database:
 
 ```bash
-# List databases
-gcloud firestore databases list --project=integration-testing-476513
+# Delete test database
+make clean-integration-db
 
-# Delete test database if it exists
+# Or use gcloud directly
 gcloud firestore databases delete --database=ds9-test --project=integration-testing-476513
 ```
 
 ### Costs
 
-Integration tests create a temporary database and a small number of entities (typically <20 per run). The database and all data are deleted after the test run completes. This should fall well within GCP free tier limits.
+Integration tests create or reuse a persistent database and a small number of entities (typically <20 per run). Entities are deleted after each test run. The persistent database incurs minimal costs and should fall well within GCP free tier limits.

datastore.go

@@ -87,7 +87,7 @@ func NewClientWithDatabase(ctx context.Context, projID, dbID string) (*Client, e
 
 	if projID == "" {
 		logger.InfoContext(ctx, "project ID not provided, fetching from metadata server")
-		pid, err := projectID(ctx)
+		pid, err := auth.ProjectID(ctx)
 		if err != nil {
 			logger.ErrorContext(ctx, "failed to get project ID from metadata server", "error", err)
 			return nil, fmt.Errorf("project ID required: %w", err)
@@ -105,11 +105,6 @@ func NewClientWithDatabase(ctx context.Context, projID, dbID string) (*Client, e
 	}, nil
 }
 
-// projectID retrieves the project ID using the auth package.
-func projectID(ctx context.Context) (string, error) {
-	return auth.ProjectID(ctx)
-}
-
 // accessToken retrieves an access token using the auth package.
 func accessToken(ctx context.Context) (string, error) {
 	return auth.AccessToken(ctx)
@@ -220,26 +215,25 @@ func doRequest(ctx context.Context, logger *slog.Logger, url string, jsonData []
 			"body_size", len(body),
 			"attempt", attempt+1)
 
+		// Success
+		if resp.StatusCode == http.StatusOK {
+			return body, nil
+		}
+
 		// Don't retry on 4xx errors (client errors)
 		if resp.StatusCode >= 400 && resp.StatusCode < 500 {
 			if resp.StatusCode == http.StatusNotFound {
 				logger.DebugContext(ctx, "entity not found", "status_code", resp.StatusCode)
 			} else {
 				logger.WarnContext(ctx, "client error", "status_code", resp.StatusCode, "body", string(body))
 			}
-			if resp.StatusCode != http.StatusOK {
-				return nil, fmt.Errorf("request failed with status %d: %s", resp.StatusCode, string(body))
-			}
-			return body, nil
+			return nil, fmt.Errorf("request failed with status %d: %s", resp.StatusCode, string(body))
 		}
 
-		// Success
+		// Unexpected 2xx/3xx status codes
 		if resp.StatusCode < 400 {
-			if resp.StatusCode != http.StatusOK {
-				logger.WarnContext(ctx, "unexpected success status code", "status_code", resp.StatusCode)
-				return nil, fmt.Errorf("unexpected status code %d: %s", resp.StatusCode, string(body))
-			}
-			return body, nil
+			logger.WarnContext(ctx, "unexpected non-200 success status", "status_code", resp.StatusCode)
+			return nil, fmt.Errorf("unexpected status code %d: %s", resp.StatusCode, string(body))
 		}
 
 		// 5xx errors - retry
@@ -616,6 +610,34 @@ func (c *Client) DeleteMulti(ctx context.Context, keys []*Key) error {
 	return nil
 }
 
+// DeleteAllByKind deletes all entities of a given kind.
+// This method queries for all keys and then deletes them in batches.
+func (c *Client) DeleteAllByKind(ctx context.Context, kind string) error {
+	c.logger.InfoContext(ctx, "deleting all entities by kind", "kind", kind)
+
+	// Query for all keys of this kind
+	q := NewQuery(kind).KeysOnly()
+	keys, err := c.AllKeys(ctx, q)
+	if err != nil {
+		c.logger.ErrorContext(ctx, "failed to query keys", "kind", kind, "error", err)
+		return fmt.Errorf("failed to query keys: %w", err)
+	}
+
+	if len(keys) == 0 {
+		c.logger.InfoContext(ctx, "no entities found to delete", "kind", kind)
+		return nil
+	}
+
+	// Delete all keys
+	if err := c.DeleteMulti(ctx, keys); err != nil {
+		c.logger.ErrorContext(ctx, "failed to delete entities", "kind", kind, "count", len(keys), "error", err)
+		return fmt.Errorf("failed to delete entities: %w", err)
+	}
+
+	c.logger.InfoContext(ctx, "deleted all entities", "kind", kind, "count", len(keys))
+	return nil
+}
+
 // keyToJSON converts a Key to its JSON representation.
 // Supports hierarchical keys with parent relationships.
 func keyToJSON(key *Key) map[string]any {
@@ -631,17 +653,17 @@ func keyToJSON(key *Key) map[string]any {
 	// Reverse to go from root to leaf
 	for i := len(keys) - 1; i >= 0; i-- {
 		k := keys[i]
-		pathElement := map[string]any{
+		elem := map[string]any{
 			"kind": k.Kind,
 		}
 
 		if k.Name != "" {
-			pathElement["name"] = k.Name
+			elem["name"] = k.Name
 		} else if k.ID != 0 {
-			pathElement["id"] = strconv.FormatInt(k.ID, 10)
+			elem["id"] = strconv.FormatInt(k.ID, 10)
 		}
 
-		path = append(path, pathElement)
+		path = append(path, elem)
 	}
 
 	return map[string]any{
@@ -996,14 +1018,18 @@ func keyFromJSON(keyData any) (*Key, error) {
 }
 
 // Transaction represents a Datastore transaction.
+// Note: This struct stores context for API compatibility with Google's official
+// cloud.google.com/go/datastore library, which uses the same pattern.
 type Transaction struct {
+	ctx       context.Context //nolint:containedctx // Required for API compatibility with cloud.google.com/go/datastore
 	client    *Client
 	id        string
 	mutations []map[string]any
 }
 
 // RunInTransaction runs a function in a transaction.
 // The function should use the transaction's Get and Put methods.
+// API compatible with cloud.google.com/go/datastore.
 func (c *Client) RunInTransaction(ctx context.Context, f func(*Transaction) error) error {
 	const maxTxRetries = 3
 	var lastErr error
@@ -1067,6 +1093,7 @@ func (c *Client) RunInTransaction(ctx context.Context, f func(*Transaction) erro
 		}
 
 		tx := &Transaction{
+			ctx:    ctx,
 			client: c,
 			id:     txResp.Transaction,
 		}
@@ -1118,16 +1145,13 @@ func (c *Client) RunInTransaction(ctx context.Context, f func(*Transaction) erro
 }
 
 // Get retrieves an entity within the transaction.
+// API compatible with cloud.google.com/go/datastore.
 func (tx *Transaction) Get(key *Key, dst any) error {
 	if key == nil {
 		return errors.New("key cannot be nil")
 	}
 
-	// For simplicity, we'll use a context.Background() here
-	// In a production implementation, you might want to pass context through
-	ctx := context.Background()
-
-	token, err := accessToken(ctx)
+	token, err := accessToken(tx.ctx)
 	if err != nil {
 		return fmt.Errorf("failed to get access token: %w", err)
 	}
@@ -1151,7 +1175,7 @@ func (tx *Transaction) Get(key *Key, dst any) error {
 	}
 
 	url := fmt.Sprintf("%s/projects/%s:lookup", apiURL, tx.client.projectID)
-	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(jsonData))
+	req, err := http.NewRequestWithContext(tx.ctx, http.MethodPost, url, bytes.NewReader(jsonData))
 	if err != nil {
 		return err
 	}