From c551d42b0a99b430c7c53690b949b3077e5b0ba9 Mon Sep 17 00:00:00 2001
From: Evan Gordon <evango@google.com>
Date: Thu, 16 May 2024 14:43:02 -0700
Subject: [PATCH] GCS testhelper can now list stored objects.

PiperOrigin-RevId: 634534468
---
 gcs/gcs_test.go    | 15 ++++++++-
 testhelpers/gcs.go | 84 ++++++++++++++++++++++++++++++++++++----------
 2 files changed, 81 insertions(+), 18 deletions(-)

diff --git a/gcs/gcs_test.go b/gcs/gcs_test.go
index bd4777a..7ae85bd 100644
--- a/gcs/gcs_test.go
+++ b/gcs/gcs_test.go
@@ -41,11 +41,24 @@ func TestGCSIsBucketInProjectId(t *testing.T) {
 	}
 	for _, tc := range cases {
 		server := testhelpers.NewGCSServer(t)
+		ctx := context.Background()
 
-		gcsClient, err := NewClient(context.Background(), tc.bucketName, server.URL())
+		gcsClient, err := NewClient(ctx, tc.bucketName, server.URL())
 		if err != nil {
 			t.Error("Unexpected error when getting NewClient: ", err)
 		}
+		if tc.wantIsBucketInProject {
+			// Add a file to the bucket so that the bucket is created.
+			writeCloser := gcsClient.GetFileWriter(ctx, "testFile")
+			_, err = writeCloser.Write([]byte("Hello World"))
+			if err != nil {
+				t.Error("Unexpected error when writing file: ", err)
+			}
+			err = writeCloser.Close()
+			if err != nil {
+				t.Error("Unexpected error when closing file and uploading data to GCS: ", err)
+			}
+		}
 
 		gotIsBucketInProject, err := gcsClient.IsBucketInProject(context.Background(), "project")
 		if err != nil {
diff --git a/testhelpers/gcs.go b/testhelpers/gcs.go
index fe78209..18a34cd 100644
--- a/testhelpers/gcs.go
+++ b/testhelpers/gcs.go
@@ -23,6 +23,8 @@ import (
 	"mime/multipart"
 	"net/http"
 	"net/http/httptest"
+	"regexp"
+	"slices"
 	"strings"
 	"sync"
 	"testing"
@@ -107,36 +109,84 @@ func (gs *GCSServer) URL() string {
 }
 
 const uploadPathPrefix = "/upload/storage/v1/b/"
-const listPathPrefix = "/b"
+
+// this should match for paths like:
+// /b - list buckets
+// /b/bucketName/o - list objects
+var listPathRegex = regexp.MustCompile(`^/b(?:/.*/o|)$`)
 
 func (gs *GCSServer) handleHTTP(w http.ResponseWriter, req *http.Request) {
 	if strings.HasPrefix(req.URL.Path, uploadPathPrefix) {
 		gs.handleUpload(w, req)
-	} else if strings.HasPrefix(req.URL.Path, listPathPrefix) {
+	} else if listPathRegex.MatchString(req.URL.Path) {
 		gs.handleList(w, req)
 	} else {
 		gs.handleDownload(w, req)
 	}
 }
 
+// A simple struct to hold the json response for the list buckets and objects calls.
+type objectIterResponse struct {
+	Kind          string
+	NextPageToken string
+	Items         []objectAttrsResponse
+}
+
+// Holds the file or bucket attributes for the list buckets / objects calls.
+type objectAttrsResponse struct {
+	Kind   string
+	ID     string
+	Name   string
+	Bucket string
+	Prefix string
+}
+
+// handleList handles the list buckets and list objects calls.
+// Does not support pagination (will only ever return a single item).
+// TODO b/341405229 - add support for arbitrary buckets.
 func (gs *GCSServer) handleList(w http.ResponseWriter, req *http.Request) {
-	r := struct {
-		Kind          string
-		NextPageToken string
-		Items         []struct {
-			Kind string
-			ID   string
-			Name string
+	var r objectIterResponse
+	if req.URL.Path == "/b" {
+		// List all buckets.
+		r = objectIterResponse{
+			Kind:  "storage#buckets",
+			Items: []objectAttrsResponse{},
+		}
+		for key := range gs.objects {
+			currBucket := objectAttrsResponse{
+				Kind: "storage#bucket",
+				ID:   key.bucket,
+				Name: key.bucket,
+			}
+			if slices.Contains(r.Items, currBucket) {
+				continue
+			}
+			r.Items = append(r.Items, currBucket)
 		}
-	}{
-		Kind:          "storage#buckets",
-		NextPageToken: "",
-		Items: []struct {
-			Kind string
-			ID   string
-			Name string
-		}{{Kind: "storage#bucket", ID: "bucketName", Name: "bucketName"}},
+	} else if strings.HasPrefix(req.URL.Path, "/b/") && strings.HasSuffix(req.URL.Path, "/o") {
+		// List all objects in a bucket.
+		bucketName := strings.Split(strings.TrimPrefix(req.URL.Path, "/b/"), "/")[0]
+		// "/b/bucketName/o"
+		r = objectIterResponse{
+			Kind:  "storage#objects",
+			Items: []objectAttrsResponse{},
+		}
+		queryPrefix := req.URL.Query().Get("prefix")
+		// find all objects in the server that match the prefix.
+		for key := range gs.objects {
+			if strings.Contains(key.name, queryPrefix) && key.bucket == bucketName {
+				r.Items = append(r.Items, objectAttrsResponse{
+					Kind:   "storage#object",
+					Name:   key.name,
+					Bucket: bucketName,
+					Prefix: queryPrefix,
+				})
+			}
+		}
+	} else {
+		gs.t.Fatalf("unrecognised list endpoint %s", req.URL.Path)
 	}
+
 	j, err := json.Marshal(r)
 	if err != nil {
 		gs.t.Fatalf("failed to marshal json in GCS handleList: %v", err)