test(firestore): Cleanup before creating new docs (#4860)

* test(firestore): Cleanup before creating new docs * test(firestore): Skip index creation if already exists
GoogleCloudPlatform · Dec 13, 2024 · 7b9a8da · 7b9a8da
1 parent b1abd27
commit 7b9a8da
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 8 deletions.
diff --git a/firestore/main_test.go b/firestore/main_test.go
@@ -26,6 +26,8 @@ import (
 	"cloud.google.com/go/firestore"
 	apiv1 "cloud.google.com/go/firestore/apiv1/admin"
 	"cloud.google.com/go/firestore/apiv1/admin/adminpb"
+	"github.com/GoogleCloudPlatform/golang-samples/internal/testutil"
+	"google.golang.org/api/iterator"
 )
 
 var projectID string
@@ -88,6 +90,9 @@ func vectorSearchSetup() func() {
 
 	cleanups := []func(){}
 
+	// Delete existing documents
+	deleteTestCollection(projectID, vectorCollName)
+
 	// Create documents
 	cleanupDocs := createCoffeeBeans(projectID, vectorCollName)
 	cleanups = append(cleanups, cleanupDocs)
@@ -249,10 +254,41 @@ func createCoffeeBeans(projectID string, collName string) func() {
 
 	return func() {
 		for _, ref := range docRefs {
-			_, err := ref.Delete(ctx)
+			testutil.RetryWithoutTest(5, 5*time.Second, func(r *testutil.R) {
+				_, err := ref.Delete(ctx)
+				if err != nil {
+					log.Printf("Error deleting document %v: %s", ref, err)
+					r.Fail()
+				}
+			})
+		}
+	}
+}
+
+func deleteTestCollection(projectID, collName string) {
+	ctx := context.Background()
+	client, err := firestore.NewClient(ctx, projectID)
+	if err != nil {
+		log.Fatalf("Failed to create client: %v", err)
+	}
+	defer client.Close()
+
+	// Delete all documents in the collName collection.
+	iter := client.Collection(collName).Documents(ctx)
+	for {
+		doc, err := iter.Next()
+		if err == iterator.Done {
+			break
+		}
+		if err != nil {
+			log.Fatalf("Failed to iterate: %v", err)
+		}
+		testutil.RetryWithoutTest(5, 5*time.Second, func(r *testutil.R) {
+			_, err = doc.Ref.Delete(ctx)
 			if err != nil {
-				log.Printf("An error has occurred: %s", err)
+				log.Printf("Error deleting document %v: %s", doc.Ref, err)
+				r.Fail()
 			}
-		}
+		})
 	}
 }
diff --git a/firestore/query_multiple_inequality_test.go b/firestore/query_multiple_inequality_test.go
@@ -25,6 +25,8 @@ import (
 	"cloud.google.com/go/firestore"
 	apiv1 "cloud.google.com/go/firestore/apiv1/admin"
 	"cloud.google.com/go/firestore/apiv1/admin/adminpb"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
 )
 
 func setupClientAndCities(t *testing.T, projectID string) (*firestore.Client, func()) {
@@ -148,15 +150,25 @@ func TestMultipleInequalitiesQuery(t *testing.T) {
 			Fields:     adminPbIndexFields,
 		},
 	}
+	var createdIndex *adminpb.Index
+	var waitErr error
 	op, createErr := adminClient.CreateIndex(ctx, req)
 	if createErr != nil {
-		t.Fatalf("CreateIndex: %v", createErr)
-	}
-	createdIndex, waitErr := op.Wait(ctx)
-	if waitErr != nil {
-		t.Fatalf("CreateIndexes failed. Wait: %v", waitErr)
+		s, ok := status.FromError(createErr)
+		if !ok || s.Code() != codes.AlreadyExists {
+			// Fail the test only if the index does not already exist
+			t.Fatalf("CreateIndex: %v", createErr)
+		}
+	} else {
+		createdIndex, waitErr = op.Wait(ctx)
+		if waitErr != nil {
+			t.Fatalf("CreateIndex failed. Wait: %v", waitErr)
+		}
 	}
 	t.Cleanup(func() {
+		if createdIndex == nil {
+			return
+		}
 		if err = adminClient.DeleteIndex(ctx, &adminpb.DeleteIndexRequest{
 			Name: createdIndex.Name,
 		}); err != nil {