Skip to content

Commit

Permalink
feat: enable more methods for retry test support in gRPC (#628)
Browse files Browse the repository at this point in the history 
* feat: enable more methods for retry test support in gRPC

* lint

* notifications test cases not applicable

* remove unneeded decorator
  • Loading branch information
@cojenco
cojenco authored Apr 15, 2024
1 parent 7b33695 commit 8917c3f
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 18 deletions.
36 changes: 24 additions & 12 deletions testbench/database.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -446,22 +446,34 @@ def __validate_injected_failure_description(self, failure):
def __validate_grpc_method_implemented_retry(self, method):
"""Returns Unimplemented 501 for methods that are not yet supported.
Temporary validation while adding Retry Test API support in gRPC."""
implemented_grpc_w_retry = {
"storage.buckets.get",
"storage.buckets.getIamPolicy",
"storage.buckets.list",
"storage.hmacKey.get",
"storage.hmacKey.list",
not_supported_grpc_w_retry = {
"storage.bucket_acl.get",
"storage.bucket_acl.list",
"storage.bucket_acl.delete",
"storage.bucket_acl.insert",
"storage.bucket_acl.patch",
"storage.bucket_acl.update",
"storage.default_object_acl.get",
"storage.default_object_acl.list",
"storage.default_object_acl.delete",
"storage.default_object_acl.insert",
"storage.default_object_acl.patch",
"storage.default_object_acl.update",
"storage.object_acl.get",
"storage.object_acl.list",
"storage.object_acl.delete",
"storage.object_acl.insert",
"storage.object_acl.patch",
"storage.object_acl.update",
"storage.notifications.delete",
"storage.notifications.get",
"storage.notifications.insert",
"storage.notifications.list",
"storage.objects.insert",
"storage.objects.list",
"storage.objects.get",
"storage.serviceaccount.get",
}
if method not in implemented_grpc_w_retry:
if method in not_supported_grpc_w_retry:
testbench.error.unimplemented(
"Retry Test API support for the requested method <%s> in GRPC" % method,
"Retry Test API not supported for the requested method <%s> in GRPC"
% method,
None,
)

Expand Down
19 changes: 15 additions & 4 deletions testbench/grpc_server.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -161,6 +161,7 @@ def __init__(self, db, echo_metadata=False):
self.db.insert_test_bucket()
self.echo_metadata = echo_metadata

@retry_test(method="storage.buckets.delete")
def DeleteBucket(self, request, context):
self.db.delete_bucket(
request.name,
Expand All @@ -169,7 +170,7 @@ def DeleteBucket(self, request, context):
)
return empty_pb2.Empty()

@retry_test("storage.buckets.get")
@retry_test(method="storage.buckets.get")
def GetBucket(self, request, context):
bucket = self.db.get_bucket(
request.name,
Expand All @@ -178,12 +179,13 @@ def GetBucket(self, request, context):
)
return bucket.metadata

@retry_test(method="storage.buckets.insert")
def CreateBucket(self, request, context):
bucket, _ = gcs.bucket.Bucket.init_grpc(request, context)
self.db.insert_bucket(bucket, context)
return bucket.metadata

@retry_test("storage.buckets.list")
@retry_test(method="storage.buckets.list")
def ListBuckets(self, request, context):
if not request.parent.startswith("projects/"):
return testbench.error.invalid(
Expand Down Expand Up @@ -221,6 +223,7 @@ def filter(bucket):
]
return storage_pb2.ListBucketsResponse(buckets=buckets)

@retry_test(method="storage.buckets.lockRetentionPolicy")
def LockBucketRetentionPolicy(self, request, context):
if request.if_metageneration_match <= 0:
return testbench.error.invalid(
Expand Down Expand Up @@ -256,11 +259,13 @@ def GetIamPolicy(self, request, context):
bucket = self.db.get_bucket(request.resource, context)
return bucket.iam_policy

@retry_test(method="storage.buckets.setIamPolicy")
def SetIamPolicy(self, request, context):
bucket = self.db.get_bucket(request.resource, context)
bucket.set_iam_policy(request, context)
return bucket.iam_policy

@retry_test(method="storage.buckets.testIamPermissions")
def TestIamPermissions(self, request, context):
# If the bucket does not exist this will return an error
_ = self.db.get_bucket(request.resource, context)
Expand All @@ -269,6 +274,7 @@ def TestIamPermissions(self, request, context):
permissions=request.permissions
)

@retry_test(method="storage.buckets.patch")
def UpdateBucket(self, request, context):
intersection = field_mask_pb2.FieldMask(
paths=[
Expand Down Expand Up @@ -377,7 +383,6 @@ def DeleteNotificationConfig(self, request, context):
bucket.delete_notification(notification_id, context)
return empty_pb2.Empty()

@retry_test(method="storage.notifications.get")
def GetNotificationConfig(self, request, context):
bucket_name, notification_id = self._decompose_notification_name(
request.name, context
Expand All @@ -404,7 +409,6 @@ def CreateNotificationConfig(self, request, context):
rest = bucket.insert_notification(json.dumps(notification), context)
return self._notification_from_rest(rest, request.parent)

@retry_test("storage.notifications.list")
def ListNotificationConfigs(self, request, context):
bucket = self.db.get_bucket(request.parent, context)
items = bucket.list_notifications(context).get("items", [])
Expand All @@ -414,6 +418,7 @@ def ListNotificationConfigs(self, request, context):
]
)

@retry_test(method="storage.objects.compose")
def ComposeObject(self, request, context):
if len(request.source_objects) == 0:
return testbench.error.missing(
Expand Down Expand Up @@ -479,6 +484,7 @@ def precondition(_, live_version, ctx):
)
return blob.metadata

@retry_test(method="storage.objects.delete")
def DeleteObject(self, request, context):
self.db.delete_object(
request.bucket,
Expand Down Expand Up @@ -571,6 +577,7 @@ def ReadObject(self, request, context):
content_range = None
start = start + size

@retry_test(method="storage.objects.patch")
def UpdateObject(self, request, context):
intersection = field_mask_pb2.FieldMask(
paths=[
Expand Down Expand Up @@ -716,6 +723,7 @@ def ListObjects(self, request, context):
items, prefixes = self.db.list_object(request, request.parent, context)
return storage_pb2.ListObjectsResponse(objects=items, prefixes=prefixes)

@retry_test(method="storage.objects.rewrite")
def RewriteObject(self, request, context):
token = request.rewrite_token
if token == "":
Expand Down Expand Up @@ -817,6 +825,7 @@ def _hmac_key_metadata_from_rest(self, rest):
rest["update_time"] = rest.pop("updated")
return json_format.ParseDict(rest, storage_pb2.HmacKeyMetadata())

@retry_test(method="storage.hmacKey.create")
def CreateHmacKey(self, request, context):
if not request.project.startswith("projects/"):
return testbench.error.invalid(
Expand All @@ -835,6 +844,7 @@ def CreateHmacKey(self, request, context):
metadata=self._hmac_key_metadata_from_rest(rest.get("metadata")),
)

@retry_test(method="storage.hmacKey.delete")
def DeleteHmacKey(self, request, context):
if not request.project.startswith("projects/"):
return testbench.error.invalid(
Expand Down Expand Up @@ -888,6 +898,7 @@ def ListHmacKeys(self, request, context):
]
)

@retry_test(method="storage.hmacKey.update")
def UpdateHmacKey(self, request, context):
if request.update_mask.paths == []:
return testbench.error.invalid(
Expand Down
4 changes: 2 additions & 2 deletions tests/test_database.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -465,11 +465,11 @@ def test_insert_retry_test_invalid_transport(self):

def test_insert_retry_test_unimplemented_grpc_method(self):
database = testbench.database.Database.init()
database.insert_supported_methods(["storage.resumable.upload"])
database.insert_supported_methods(["storage.bucket_acl.get"])

with self.assertRaises(testbench.error.RestException) as rest:
_ = database.insert_retry_test(
{"storage.resumable.upload": ["return-429"]}, transport="GRPC"
{"storage.bucket_acl.get": ["return-429"]}, transport="GRPC"
)
self.assertEqual(rest.exception.code, 501)

Expand Down

0 comments on commit 8917c3f

Please sign in to comment.