fix: add prefix support to list_buckets (#316)

* fix: add prefix support to list_buckets

* fix formatting

testbench/database.py

@@ -97,11 +97,20 @@ def insert_bucket(self, bucket, context):
             self._objects[bucket.metadata.name] = {}
             self._live_generations[bucket.metadata.name] = {}
 
-    def list_bucket(self, project_id, context):
+    def list_bucket(self, project_id, prefix, context):
         with self._resources_lock:
             if project_id is None or project_id.endswith("-"):
                 testbench.error.invalid("Project id %s" % project_id, context)
-            return self._buckets.values()
+            if not prefix:
+                return self._buckets.values()
+
+            buckets = []
+            for bucket in self._buckets.values():
+                name = bucket.metadata.name
+                if name.find(prefix) == 0:
+                    buckets.append(bucket)
+
+            return buckets
 
     def delete_bucket(self, bucket_name, context, preconditions=[]):
         with self._resources_lock:

testbench/grpc_server.py

@@ -66,6 +66,10 @@ def ListBuckets(self, request, context):
                 "invalid format for parent=%s" % request.parent, context
             )
         project = request.parent[len("projects/") :]
+        prefix = request.prefix
+        if prefix:
+            prefix = "projects/_/buckets/" + prefix
+
         if len(request.read_mask.paths) == 0:
             # By default we need to filter out `acl`, `default_object_acl`, and `owner`
             def filter(bucket):
@@ -90,7 +94,9 @@ def filter(bucket):
                 request.read_mask.MergeMessage(bucket, b)
                 return b
 
-        buckets = [filter(b.metadata) for b in self.db.list_bucket(project, context)]
+        buckets = [
+            filter(b.metadata) for b in self.db.list_bucket(project, prefix, context)
+        ]
         return storage_pb2.ListBucketsResponse(buckets=buckets)
 
     def LockBucketRetentionPolicy(self, request, context):

testbench/rest_server.py

@@ -190,9 +190,10 @@ def bucket_list():
     project = flask.request.args.get("project")
     projection = flask.request.args.get("projection", "noAcl")
     fields = flask.request.args.get("fields", None)
+    prefix = flask.request.args.get("prefix", "")
     response = {
         "kind": "storage#buckets",
-        "items": [bucket.rest() for bucket in db.list_bucket(project, None)],
+        "items": [bucket.rest() for bucket in db.list_bucket(project, prefix, None)],
     }
     return testbench.common.filter_response_rest(response, projection, fields)
 

tests/test_database.py

@@ -45,11 +45,16 @@ def test_bucket_crud(self):
 
         get_result = database.get_bucket("bucket-name", None)
         self.assertEqual(bucket.metadata, get_result.metadata)
-        list_result = database.list_bucket("test-project-id", None)
+        list_result = database.list_bucket("test-project-id", "", None)
         names = {b.metadata.bucket_id for b in list_result}
         self.assertEqual(names, {"bucket-name"})
+        list_result = database.list_bucket(
+            "test-project-id", "nonexistent-prefix", None
+        )
+        names = {b.metadata.name for b in list_result}
+        self.assertEqual(names, set())
         database.delete_bucket("bucket-name", None)
-        list_result = database.list_bucket("test-project-id", None)
+        list_result = database.list_bucket("test-project-id", "", None)
         names = {b.metadata.name for b in list_result}
         self.assertEqual(names, set())
 
@@ -91,7 +96,7 @@ def test_list_bucket_invalid(self):
                 args={},
                 data=json.dumps({}),
             )
-            database.list_bucket("invalid-project-id-", None)
+            database.list_bucket("invalid-project-id-", "", None)
         self.assertEqual(rest.exception.code, 400)
 
     def test_delete_not_empty(self):
@@ -123,7 +128,7 @@ def test_insert_test_bucket(self):
         )
         os.environ.pop("GOOGLE_CLOUD_CPP_STORAGE_TEST_BUCKET_NAME", None)
         database.insert_test_bucket()
-        names = {b.metadata.name for b in database.list_bucket("", None)}
+        names = {b.metadata.name for b in database.list_bucket("", "", None)}
         self.assertEqual(names, set())
 
         os.environ["GOOGLE_CLOUD_CPP_STORAGE_TEST_BUCKET_NAME"] = "test-bucket-1"

tests/test_grpc_server.py

@@ -46,13 +46,13 @@ def test_insert_test_bucket(self):
         os.environ.pop("GOOGLE_CLOUD_CPP_STORAGE_TEST_BUCKET_NAME", None)
         database = testbench.database.Database.init()
         server = testbench.grpc_server.StorageServicer(database)
-        names = {b.metadata.name for b in database.list_bucket("", None)}
+        names = {b.metadata.name for b in database.list_bucket("", "", None)}
         self.assertEqual(names, set())
 
         os.environ["GOOGLE_CLOUD_CPP_STORAGE_TEST_BUCKET_NAME"] = "test-bucket-1"
         database = testbench.database.Database.init()
         server = testbench.grpc_server.StorageServicer(database)
-        names = {b.metadata.name for b in database.list_bucket("", None)}
+        names = {b.metadata.name for b in database.list_bucket("", "", None)}
         self.assertIn("projects/_/buckets/test-bucket-1", names)
 
     def test_delete_bucket(self):
@@ -224,6 +224,35 @@ def test_list_buckets_filter(self):
             self.assertTrue(b.HasField("owner"), msg=b)
             self.assertNotEqual(len(b.acl), 0, msg=b)
 
+    def test_list_buckets_prefix(self):
+        ids = ["a-bucket-3", "b-bucket-2", "a-bucket-1"]
+        for id in ids:
+            context = unittest.mock.Mock()
+            response = self.grpc.CreateBucket(
+                storage_pb2.CreateBucketRequest(
+                    parent="projects/test-project",
+                    bucket_id=id,
+                    bucket=storage_pb2.Bucket(),
+                ),
+                context,
+            )
+            context.abort.assert_not_called()
+            self.assertEqual(response.bucket_id, id)
+        context = unittest.mock.Mock()
+        response = self.grpc.ListBuckets(
+            storage_pb2.ListBucketsRequest(parent="projects/test-project", prefix="a-"),
+            context,
+        )
+        context.assert_not_called()
+        want = [ids[0], ids[2]]
+        expected = {("projects/_/buckets/" + id) for id in want}
+        actual = {b.name for b in response.buckets}
+        self.assertEqual(actual, actual | expected)
+        for b in response.buckets:
+            self.assertFalse(b.HasField("owner"), msg=b)
+            self.assertEqual(len(b.acl), 0, msg=b)
+            self.assertEqual(len(b.default_object_acl), 0, msg=b)
+
     def test_list_buckets_failure(self):
         context = unittest.mock.Mock()
         _ = self.grpc.ListBuckets(