Repository: googleapis/storage-testbench

Issue:
Support stalls for resumable write
Is your feature request related to a problem? Please describe.
I would like to inject a retriable stall for writes. This will be beneficial to test intermittent stall scenario.

Currently, we have support for [stall-always](https://github.com/googleapis/storage-testbench?tab=readme-ov-file#stall-always), which stalls the request every time and not useful to test the intermittent stall for resumble writes.

Describe the solution you'd like
Support injection of retriable stall using x-retry-test-id header. Also, useful if we can provide the controllable stall delay.

Describe alternatives you've considered
NA

Additional context
NA

Hints from issue comments:

Solution patch:
diff --git a/README.md b/README.md
index f498466..9e56968 100644
--- a/README.md
+++ b/README.md
@@ -252,7 +252,7 @@ curl -H "x-retry-test-id: 1d05c20627844214a9ff7cbcf696317d" "http://localhost:91
 | return-broken-stream                      | [HTTP] Testbench will fail after a few downloaded bytes <br> [GRPC] Testbench will fail with `UNAVAILABLE` after a few downloaded bytes
 | return-broken-stream-after-YK             | [HTTP] Testbench will fail after YKiB of downloaded data <br> [GRPC] Testbench will fail with `UNAVAILABLE` after YKiB of downloaded data
 | return-reset-connection                   | [HTTP] Testbench will fail with a reset connection <br> [GRPC] Testbench will fail the RPC with `UNAVAILABLE`
-| stall-for-Ts-after-YK                     | [HTTP] Testbench will stall for T second after reading YKiB of downloaded data, e.g. stall-for-10s-after-12K stalls after reading 12KiB of data <br> [GRPC] Not supported
+| stall-for-Ts-after-YK                     | [HTTP] Testbench will stall for T second after reading YKiB of downloaded/uploaded data, e.g. stall-for-10s-after-12K stalls after reading/writing 12KiB of data <br> [GRPC] Not supported
 
 ## Releasing the testbench
 
@@ -264,4 +264,4 @@ Steps:
 1. Title "v0.x.x"
 1. Click Generate release notes
 1. Make sure "Set as the latest release" is checked
-1. Click "Publish Release" to release
+1. Click "Publish Release" to release
\ No newline at end of file
diff --git a/setup.py b/setup.py
index a3e34a0..ba5eed7 100644
--- a/setup.py
+++ b/setup.py
@@ -43,7 +43,7 @@ setuptools.setup(
     ],
     python_requires=">=3.8",
     install_requires=[
-        "grpcio==1.66.1",
+        "grpcio==1.67.0",
         "googleapis-common-protos==1.65.0",
         "protobuf==5.27.2",
         "flask==3.0.3",
diff --git a/testbench/common.py b/testbench/common.py
index 678a5d8..5f7399c 100644
--- a/testbench/common.py
+++ b/testbench/common.py
@@ -844,11 +844,17 @@ def handle_retry_test_instruction(database, request, socket_closer, method):
     broken_stream_after_bytes = (
         testbench.common.retry_return_broken_stream_after_bytes.match(next_instruction)
     )
+    if broken_stream_after_bytes and method == "storage.objects.get":
+        items = list(broken_stream_after_bytes.groups())
+        after_bytes = int(items[0]) * 1024
+        return __get_streamer_response_fn(
+            database, method, socket_closer, test_id, limit=after_bytes
+        )
 
     retry_stall_after_bytes_matches = testbench.common.retry_stall_after_bytes.match(
         next_instruction
     )
-    if retry_stall_after_bytes_matches:
+    if retry_stall_after_bytes_matches and method != "storage.objects.insert":
         items = list(retry_stall_after_bytes_matches.groups())
         stall_time = int(items[0])
         after_bytes = int(items[1]) * 1024
@@ -856,12 +862,6 @@ def handle_retry_test_instruction(database, request, socket_closer, method):
             database, method, test_id, limit=after_bytes, stall_time_sec=stall_time
         )
 
-    if broken_stream_after_bytes and method == "storage.objects.get":
-        items = list(broken_stream_after_bytes.groups())
-        after_bytes = int(items[0]) * 1024
-        return __get_streamer_response_fn(
-            database, method, socket_closer, test_id, limit=after_bytes
-        )
     retry_return_short_response = testbench.common.retry_return_short_response.match(
         next_instruction
     )
@@ -895,6 +895,30 @@ def gen_retry_test_decorator(db):
     return retry_test
 
 
+def get_stall_uploads_after_bytes(database, request, context=None, transport="HTTP"):
+    """Retrieve stall time and #bytes corresponding to uploads from retry test instructions."""
+    method = "storage.objects.insert"
+    test_id = request.headers.get("x-retry-test-id", None)
+    if not test_id:
+        return 0, 0, ""
+    next_instruction = None
+    if database.has_instructions_retry_test(test_id, method, transport=transport):
+        next_instruction = database.peek_next_instruction(test_id, method)
+    if not next_instruction:
+        return 0, 0, ""
+
+    stall_after_byte_matches = testbench.common.retry_stall_after_bytes.match(
+        next_instruction
+    )
+    if stall_after_byte_matches:
+        items = list(stall_after_byte_matches.groups())
+        stall_time = int(items[0])
+        after_bytes = int(items[1]) * 1024
+        return stall_time, after_bytes, test_id
+
+    return 0, 0, ""
+
+
 def get_retry_uploads_error_after_bytes(
     database, request, context=None, transport="HTTP"
 ):
@@ -919,9 +943,29 @@ def get_retry_uploads_error_after_bytes(
         error_code = int(items[0])
         after_bytes = int(items[1]) * 1024
         return error_code, after_bytes, test_id
+
     return 0, 0, ""
 
 
+def handle_stall_uploads_after_bytes(
+    upload,
+    data,
+    database,
+    stall_time,
+    after_bytes,
+    test_id=0,
+):
+    """
+    Handle stall-after-bytes instructions for resumable uploads.
+    Stall happens after given value of bytes.
+    e.g. We are uploading 120K of data then, stall-2s-after-100K will stall the request.
+    """
+    if len(upload.media) <= after_bytes and len(upload.media) + len(data) > after_bytes:
+        if test_id:
+            database.dequeue_next_instruction(test_id, "storage.objects.insert")
+        time.sleep(stall_time)
+
+
 def handle_retry_uploads_error_after_bytes(
     upload,
     data,
diff --git a/testbench/database.py b/testbench/database.py
index 30cd028..65412f1 100644
--- a/testbench/database.py
+++ b/testbench/database.py
@@ -606,6 +606,11 @@ class Database:
             "storage.default_object_acl.insert",
             "storage.default_object_acl.patch",
             "storage.default_object_acl.update",
+            "storage.hmacKey.create",
+            "storage.hmacKey.delete",
+            "storage.hmacKey.get",
+            "storage.hmacKey.list",
+            "storage.hmacKey.update",
             "storage.object_acl.get",
             "storage.object_acl.list",
             "storage.object_acl.delete",
@@ -616,6 +621,7 @@ class Database:
             "storage.notifications.get",
             "storage.notifications.insert",
             "storage.notifications.list",
+            "storage.serviceaccount.get",
         }
         if method in not_supported_grpc_w_retry:
             testbench.error.unimplemented(
diff --git a/testbench/grpc_server.py b/testbench/grpc_server.py
index 6087ef3..4a2370e 100644
--- a/testbench/grpc_server.py
+++ b/testbench/grpc_server.py
@@ -352,72 +352,6 @@ class StorageServicer(storage_pb2_grpc.StorageServicer):
         bucket.metadata.update_time.FromDatetime(now)
         return bucket.metadata
 
-    def _notification_from_rest(self, rest, bucket_name):
-        # We need to make a copy before changing any values
-        rest = rest.copy()
-        rest.pop("kind")
-        rest["name"] = bucket_name + "/notificationConfigs/" + rest.pop("id")
-        rest["topic"] = "//pubsub.googleapis.com/" + rest["topic"]
-        return json_format.ParseDict(rest, storage_pb2.NotificationConfig())
-
-    def _decompose_notification_name(self, notification_name, context):
-        loc = notification_name.find("/notificationConfigs/")
-        if loc == -1:
-            testbench.error.invalid(
-                "GetNotificationConfig() malformed notification name [%s]"
-                % notification_name,
-                context,
-            )
-            return (None, None)
-        bucket_name = notification_name[:loc]
-        notification_id = notification_name[loc + len("/notificationConfigs/") :]
-        return (bucket_name, notification_id)
-
-    def DeleteNotificationConfig(self, request, context):
-        bucket_name, notification_id = self._decompose_notification_name(
-            request.name, context
-        )
-        if bucket_name is None:
-            return None
-        bucket = self.db.get_bucket(bucket_name, context)
-        bucket.delete_notification(notification_id, context)
-        return empty_pb2.Empty()
-
-    def GetNotificationConfig(self, request, context):
-        bucket_name, notification_id = self._decompose_notification_name(
-            request.name, context
-        )
-        if bucket_name is None:
-            return None
-        bucket = self.db.get_bucket(bucket_name, context)
-        rest = bucket.get_notification(notification_id, context)
-        return self._notification_from_rest(rest, bucket_name)
-
-    def CreateNotificationConfig(self, request, context):
-        pattern = "^//pubsub.googleapis.com/projects/[^/]+/topics/[^/]+$"
-        if re.match(pattern, request.notification_config.topic) is None:
-            return testbench.error.invalid(
-                "topic names must be in"
-                + " //pubsub.googleapis.com/projects/{project-identifier}/topics/{my-topic}"
-                + " format, got=%s" % request.notification_config.topic,
-                context,
-            )
-        bucket = self.db.get_bucket(request.parent, context)
-        notification = json_format.MessageToDict(request.notification_config)
-        # Convert topic names to REST format
-        notification["topic"] = notification["topic"][len("//pubsub.googleapis.com/") :]
-        rest = bucket.insert_notification(json.dumps(notification), context)
-        return self._notification_from_rest(rest, request.parent)
-
-    def ListNotificationConfigs(self, request, context):
-        bucket = self.db.get_bucket(request.parent, context)
-        items = bucket.list_notifications(context).get("items", [])
-        return storage_pb2.ListNotificationConfigsResponse(
-            notification_configs=[
-                self._notification_from_rest(r, request.parent) for r in items
-            ]
-        )
-
     @retry_test(method="storage.objects.compose")
     def ComposeObject(self, request, context):
         if len(request.source_objects) == 0:
@@ -815,123 +749,6 @@ class StorageServicer(storage_pb2_grpc.StorageServicer):
             return storage_pb2.QueryWriteStatusResponse(resource=upload.blob.metadata)
         return storage_pb2.QueryWriteStatusResponse(persisted_size=len(upload.media))
 
-    @retry_test("storage.serviceaccount.get")
-    def GetServiceAccount(self, request, context):
-        if not request.project.startswith("projects/"):
-            return testbench.error.invalid(
-                "project name must start with projects/, got=%s" % request.project,
-                context,
-            )
-        project_id = request.project[len("projects/") :]
-        project = self.db.get_project(project_id)
-        return storage_pb2.ServiceAccount(email_address=project.service_account_email())
-
-    def _hmac_key_metadata_from_rest(self, rest):
-        rest = rest.copy()
-        rest.pop("kind", None)
-        rest["project"] = "projects/" + rest.pop("projectId")
-        rest["create_time"] = rest.pop("timeCreated")
-        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(
-                "project name must start with projects/, got=%s" % request.project,
-                context,
-            )
-        if request.service_account_email == "":
-            return testbench.error.invalid(
-                "service account email must be non-empty", context
-            )
-        project_id = request.project[len("projects/") :]
-        project = self.db.get_project(project_id)
-        rest = project.insert_hmac_key(request.service_account_email)
-        return storage_pb2.CreateHmacKeyResponse(
-            secret_key_bytes=base64.b64decode(rest.get("secret").encode("utf-8")),
-            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(
-                "project name must start with projects/, got=%s" % request.project,
-                context,
-            )
-        project_id = request.project[len("projects/") :]
-        project = self.db.get_project(project_id)
-        project.delete_hmac_key(request.access_id, context)
-        return empty_pb2.Empty()
-
-    @retry_test("storage.hmacKey.get")
-    def GetHmacKey(self, request, context):
-        if not request.project.startswith("projects/"):
-            return testbench.error.invalid(
-                "project name must start with projects/, got=%s" % request.project,
-                context,
-            )
-        project_id = request.project[len("projects/") :]
-        project = self.db.get_project(project_id)
-        rest = project.get_hmac_key(request.access_id, context)
-        return self._hmac_key_metadata_from_rest(rest)
-
-    @retry_test("storage.hmacKey.list")
-    def ListHmacKeys(self, request, context):
-        if not request.project.startswith("projects/"):
-            return testbench.error.invalid(
-                "project name must start with projects/, got=%s" % request.project,
-                context,
-            )
-        project_id = request.project[len("projects/") :]
-        project = self.db.get_project(project_id)
-
-        items = []
-        sa_email = request.service_account_email
-        if len(sa_email) != 0:
-            service_account = project.service_account(sa_email)
-            if service_account:
-                items = service_account.key_items()
-        else:
-            for sa in project.service_accounts.values():
-                items.extend(sa.key_items())
-
-        state_filter = lambda x: x.get("state") != "DELETED"
-        if request.show_deleted_keys:
-            state_filter = lambda x: True
-
-        return storage_pb2.ListHmacKeysResponse(
-            hmac_keys=[
-                self._hmac_key_metadata_from_rest(i) for i in items if state_filter(i)
-            ]
-        )
-
-    @retry_test(method="storage.hmacKey.update")
-    def UpdateHmacKey(self, request, context):
-        if request.update_mask.paths == []:
-            return testbench.error.invalid(
-                "UpdateHmacKey() with an empty update mask", context
-            )
-        if request.update_mask.paths != ["state"]:
-            return testbench.error.invalid(
-                "UpdateHmacKey() only the `state` field can be modified [%s]"
-                % ",".join(request.update_mask.paths),
-                context,
-            )
-        project_id = request.hmac_key.project
-        if not project_id.startswith("projects/"):
-            return testbench.error.invalid(
-                "project name must start with projects/, got=%s" % project_id, context
-            )
-        project_id = project_id[len("projects/") :]
-        project = self.db.get_project(project_id)
-        payload = {"state": request.hmac_key.state}
-        if request.hmac_key.etag != "":
-            payload["etag"] = request.hmac_key.etag
-        rest = project.update_hmac_key(request.hmac_key.access_id, payload, context)
-        return self._hmac_key_metadata_from_rest(rest)
-
 
 def run(port, database, echo_metadata=False):
     server = grpc.server(futures.ThreadPoolExecutor(max_workers=1))
diff --git a/testbench/rest_server.py b/testbench/rest_server.py
index 9630840..72a0ae8 100644
--- a/testbench/rest_server.py
+++ b/testbench/rest_server.py
@@ -16,6 +16,7 @@ import argparse
 import datetime
 import json
 import logging
+import time
 
 import flask
 from google.protobuf import json_format
@@ -978,6 +979,18 @@ def object_insert(bucket_name):
         blob, projection = gcs_type.object.Object.init_media(flask.request, bucket)
     elif upload_type == "multipart":
         blob, projection = gcs_type.object.Object.init_multipart(flask.request, bucket)
+        # Handle stall for full uploads.
+        testbench.common.extract_instruction(request, context=None)
+        (
+            stall_time,
+            after_bytes,
+            test_id,
+        ) = testbench.common.get_stall_uploads_after_bytes(db, request)
+        if stall_time and len(blob.media) >= after_bytes:
+            if test_id:
+                db.dequeue_next_instruction(test_id, "storage.objects.insert")
+            time.sleep(stall_time)
+
     db.insert_object(
         bucket_name,
         blob,
@@ -1104,6 +1117,23 @@ def resumable_upload_chunk(bucket_name):
                 chunk_last_byte,
                 test_id,
             )
+
+        testbench.common.extract_instruction(request, context=None)
+        (
+            stall_time,
+            after_bytes,
+            test_id,
+        ) = testbench.common.get_stall_uploads_after_bytes(db, request)
+
+        if stall_time:
+            testbench.common.handle_stall_uploads_after_bytes(
+                upload,
+                data,
+                db,
+                stall_time,
+                after_bytes,
+                test_id,
+            )
         # The testbench should ignore any request bytes that have already been persisted,
         # to be aligned with GCS behavior (https://cloud.google.com/storage/docs/resumable-uploads#resent-data).
         # Thus we validate chunk_first_byte against last_byte_persisted.


Test patch:
diff --git a/tests/test_grpc_server.py b/tests/test_grpc_server.py
index 3c17883..bd4e0a4 100755
--- a/tests/test_grpc_server.py
+++ b/tests/test_grpc_server.py
@@ -638,180 +638,6 @@ class TestGrpc(unittest.TestCase):
             grpc.StatusCode.INVALID_ARGUMENT, unittest.mock.ANY
         )
 
-    def test_delete_notification(self):
-        context = unittest.mock.Mock()
-        create = self.grpc.CreateNotificationConfig(
-            storage_pb2.CreateNotificationConfigRequest(
-                parent="projects/_/buckets/bucket-name",
-                notification_config=storage_pb2.NotificationConfig(
-                    topic="//pubsub.googleapis.com/projects/test-project-id/topics/test-topic",
-                    custom_attributes={"key": "value"},
-                    payload_format="JSON_API_V1",
-                ),
-            ),
-            context,
-        )
-        self.assertTrue(
-            create.name.startswith(
-                "projects/_/buckets/bucket-name/notificationConfigs/"
-            ),
-            msg=create,
-        )
-
-        context = unittest.mock.Mock()
-        _ = self.grpc.DeleteNotificationConfig(
-            storage_pb2.GetNotificationConfigRequest(name=create.name), context
-        )
-        context.abort.assert_not_called()
-
-        # The code depends on `context.abort()` raising an exception.
-        context = unittest.mock.Mock()
-        context.abort = unittest.mock.MagicMock()
-        context.abort.side_effect = grpc.RpcError()
-        with self.assertRaises(grpc.RpcError):
-            _ = self.grpc.GetNotificationConfig(
-                storage_pb2.GetNotificationConfigRequest(name=create.name), context
-            )
-        context.abort.assert_called_once_with(
-            grpc.StatusCode.NOT_FOUND, unittest.mock.ANY
-        )
-
-    def test_delete_notification_malformed(self):
-        context = unittest.mock.Mock()
-        _ = self.grpc.DeleteNotificationConfig(
-            storage_pb2.GetNotificationConfigRequest(name=""), context
-        )
-        context.abort.assert_called_once_with(
-            grpc.StatusCode.INVALID_ARGUMENT, unittest.mock.ANY
-        )
-
-    def test_get_notification(self):
-        context = unittest.mock.Mock()
-        topic_name = (
-            "//pubsub.googleapis.com/projects/test-project-id/topics/test-topic"
-        )
-        response = self.grpc.CreateNotificationConfig(
-            storage_pb2.CreateNotificationConfigRequest(
-                parent="projects/_/buckets/bucket-name",
-                notification_config=storage_pb2.NotificationConfig(
-                    topic=topic_name,
-                    custom_attributes={"key": "value"},
-                    payload_format="JSON_API_V1",
-                ),
-            ),
-            context,
-        )
-        self.assertTrue(
-            response.name.startswith(
-                "projects/_/buckets/bucket-name/notificationConfigs/"
-            ),
-            msg=response,
-        )
-
-        context = unittest.mock.Mock()
-        get = self.grpc.GetNotificationConfig(
-            storage_pb2.GetNotificationConfigRequest(name=response.name), context
-        )
-        self.assertEqual(get, response)
-        self.assertEqual(get.topic, topic_name)
-
-    def test_get_notification_malformed(self):
-        context = unittest.mock.Mock()
-        _ = self.grpc.GetNotificationConfig(
-            storage_pb2.GetNotificationConfigRequest(
-                name="projects/_/buckets/bucket-name/"
-            ),
-            context,
-        )
-        context.abort.assert_called_once_with(
-            grpc.StatusCode.INVALID_ARGUMENT, unittest.mock.ANY
-        )
-
-    def test_create_notification(self):
-        topic_name = (
-            "//pubsub.googleapis.com/projects/test-project-id/topics/test-topic"
-        )
-        context = unittest.mock.Mock()
-        response = self.grpc.CreateNotificationConfig(
-            storage_pb2.CreateNotificationConfigRequest(
-                parent="projects/_/buckets/bucket-name",
-                notification_config=storage_pb2.NotificationConfig(
-                    topic=topic_name,
-                    custom_attributes={"key": "value"},
-                    payload_format="JSON_API_V1",
-                ),
-            ),
-            context,
-        )
-        self.assertTrue(
-            response.name.startswith(
-                "projects/_/buckets/bucket-name/notificationConfigs/"
-            ),
-            msg=response,
-        )
-        self.assertEqual(response.topic, topic_name)
-
-        # Invalid topic names return an error
-        invalid_topic_names = [
-            "",
-            "//pubsub.googleapis.com",
-            "//pubsub.googleapis.com/",
-            "//pubsub.googleapis.com/projects",
-            "//pubsub.googleapis.com/projects/",
-            "//pubsub.googleapis.com/projects/test-project-id",
-            "//pubsub.googleapis.com/projects/test-project-id/",
-            "//pubsub.googleapis.com/projects/test-project-id/topics",
-            "//pubsub.googleapis.com/projects/test-project-id/topics/",
-        ]
-        for topic in invalid_topic_names:
-            context = unittest.mock.Mock()
-            context.abort = unittest.mock.MagicMock()
-            context.abort.side_effect = grpc.RpcError()
-            with self.assertRaises(grpc.RpcError):
-                _ = self.grpc.CreateNotificationConfig(
-                    storage_pb2.CreateNotificationConfigRequest(
-                        parent="projects/_/buckets/bucket-name",
-                        notification_config=storage_pb2.NotificationConfig(
-                            topic=topic, payload_format="JSON_API_V1"
-                        ),
-                    ),
-                    context,
-                )
-            context.abort.assert_called_once_with(
-                grpc.StatusCode.INVALID_ARGUMENT, unittest.mock.ANY
-            )
-
-    def test_list_notifications(self):
-        expected = set()
-        topics = [
-            "//pubsub.googleapis.com/projects/test-project-id/topics/test-topic-1",
-            "//pubsub.googleapis.com/projects/test-project-id/topics/test-topic-2",
-        ]
-        for topic in topics:
-            context = unittest.mock.Mock()
-            response = self.grpc.CreateNotificationConfig(
-                storage_pb2.CreateNotificationConfigRequest(
-                    parent="projects/_/buckets/bucket-name",
-                    notification_config=storage_pb2.NotificationConfig(
-                        topic=topic,
-                        custom_attributes={"key": "value"},
-                        payload_format="JSON_API_V1",
-                    ),
-                ),
-                context,
-            )
-            expected.add(response.name)
-
-        context = unittest.mock.Mock()
-        response = self.grpc.ListNotificationConfigs(
-            storage_pb2.ListNotificationConfigsRequest(
-                parent="projects/_/buckets/bucket-name"
-            ),
-            context,
-        )
-        names = {n.name for n in response.notification_configs}
-        self.assertEqual(names, expected)
-
     def test_compose_object(self):
         payloads = {
             "fox": b"The quick brown fox jumps over the lazy dog\n",
@@ -1867,298 +1693,6 @@ class TestGrpc(unittest.TestCase):
             response_names = [o.name for o in response.objects]
             self.assertEqual(response_names, case["expected"], msg=case)
 
-    def test_get_service_account(self):
-        context = unittest.mock.Mock()
-        response = self.grpc.GetServiceAccount(
-            storage_pb2.GetServiceAccountRequest(project="projects/test-project"),
-            context,
-        )
-        self.assertIn("@", response.email_address)
-
-    def test_get_service_account_failure(self):
-        context = unittest.mock.Mock()
-        _ = self.grpc.GetServiceAccount(
-            storage_pb2.GetServiceAccountRequest(project="invalid-format"),
-            context,
-        )
-        context.abort.assert_called_once_with(
-            grpc.StatusCode.INVALID_ARGUMENT, unittest.mock.ANY
-        )
-
-    def test_create_hmac_key(self):
-        sa_email = "test-sa@test-project-id.iam.gserviceaccount.com"
-        context = unittest.mock.Mock()
-        response = self.grpc.CreateHmacKey(
-            storage_pb2.CreateHmacKeyRequest(
-                project="projects/test-project-id",
-                service_account_email=sa_email,
-            ),
-            context,
-        )
-        self.assertNotEqual(response.secret_key_bytes, "")
-        self.assertNotEqual(response.metadata.id, "")
-        self.assertNotEqual(response.metadata.access_id, "")
-        self.assertEqual(response.metadata.service_account_email, sa_email)
-        self.assertNotEqual(response.metadata.etag, "")
-
-        context = unittest.mock.Mock()
-        _ = self.grpc.CreateHmacKey(
-            storage_pb2.CreateHmacKeyRequest(
-                project="invalid-project-name-format",
-                service_account_email=sa_email,
-            ),
-            context,
-        )
-        context.abort.assert_called_once_with(
-            grpc.StatusCode.INVALID_ARGUMENT, unittest.mock.ANY
-        )
-
-        # Missing service account emails should fail
-        context = unittest.mock.Mock()
-        _ = self.grpc.CreateHmacKey(
-            storage_pb2.CreateHmacKeyRequest(
-                project="projects/test-project-id",
-            ),
-            context,
-        )
-        context.abort.assert_called_once_with(
-            grpc.StatusCode.INVALID_ARGUMENT, unittest.mock.ANY
-        )
-
-    def test_delete_hmac_key(self):
-        sa_email = "test-sa@test-project-id.iam.gserviceaccount.com"
-        context = unittest.mock.Mock()
-        create = self.grpc.CreateHmacKey(
-            storage_pb2.CreateHmacKeyRequest(
-                project="projects/test-project-id",
-                service_account_email=sa_email,
-            ),
-            context,
-        )
-
-        metadata = storage_pb2.HmacKeyMetadata()
-        metadata.CopyFrom(create.metadata)
-        metadata.state = "INACTIVE"
-
-        context = unittest.mock.Mock()
-        response = self.grpc.UpdateHmacKey(
-            storage_pb2.UpdateHmacKeyRequest(
-                hmac_key=metadata, update_mask=field_mask_pb2.FieldMask(paths=["state"])
-            ),
-            context,
-        )
-        self.assertEqual(response.state, "INACTIVE")
-
-        context = unittest.mock.Mock()
-        _ = self.grpc.DeleteHmacKey(
-            storage_pb2.DeleteHmacKeyRequest(
-                access_id=metadata.access_id, project=metadata.project
-            ),
-            context,
-        )
-        context.abort.assert_not_called()
-
-        # The code depends on `context.abort()` raising an exception.
-        context.abort = unittest.mock.MagicMock()
-        context.abort.side_effect = grpc.RpcError()
-        with self.assertRaises(grpc.RpcError):
-            _ = self.grpc.GetHmacKey(
-                storage_pb2.GetHmacKeyRequest(
-                    access_id=metadata.access_id, project="projects/test-project-id"
-                ),
-                context,
-            )
-        context.abort.assert_called_once_with(
-            grpc.StatusCode.NOT_FOUND, unittest.mock.ANY
-        )
-
-        # Missing or malformed project id is an error
-        context = unittest.mock.Mock()
-        _ = self.grpc.DeleteHmacKey(
-            storage_pb2.DeleteHmacKeyRequest(access_id=metadata.access_id),
-            context,
-        )
-        context.abort.assert_called_once_with(
-            grpc.StatusCode.INVALID_ARGUMENT, unittest.mock.ANY
-        )
-
-    def test_get_hmac_key(self):
-        sa_email = "test-sa@test-project-id.iam.gserviceaccount.com"
-        context = unittest.mock.Mock()
-        create = self.grpc.CreateHmacKey(
-            storage_pb2.CreateHmacKeyRequest(
-                project="projects/test-project-id",
-                service_account_email=sa_email,
-            ),
-            context,
-        )
-
-        context = unittest.mock.Mock()
-        response = self.grpc.GetHmacKey(
-            storage_pb2.GetHmacKeyRequest(
-                access_id=create.metadata.access_id, project="projects/test-project-id"
-            ),
-            context,
-        )
-        self.assertEqual(response, create.metadata)
-
-        # Missing or malformed project id is an error
-        context = unittest.mock.Mock()
-        _ = self.grpc.GetHmacKey(
-            storage_pb2.GetHmacKeyRequest(access_id=create.metadata.access_id),
-            context,
-        )
-        context.abort.assert_called_once_with(
-            grpc.StatusCode.INVALID_ARGUMENT, unittest.mock.ANY
-        )
-
-    def test_list_hmac_keys(self):
-        # Create several keys for two different projects.
-        expected = []
-        sa_emails = [
-            "test-sa-1@test-project-id.iam.gserviceaccount.com",
-            "test-sa-2@test-project-id.iam.gserviceaccount.com",
-        ]
-        for sa in sa_emails:
-            for _ in [1, 2]:
-                context = unittest.mock.Mock()
-                create = self.grpc.CreateHmacKey(
-                    storage_pb2.CreateHmacKeyRequest(
-                        project="projects/test-project-id",
-                        service_account_email=sa,
-                    ),
-                    context,
-                )
-                expected.append(create.metadata)
-
-        # First test without any filtering
-        context = unittest.mock.Mock()
-        response = self.grpc.ListHmacKeys(
-            storage_pb2.ListHmacKeysRequest(project="projects/test-project-id"),
-            context,
-        )
-        expected_access_ids = {k.access_id for k in expected}
-        self.assertEqual({k.access_id for k in response.hmac_keys}, expected_access_ids)
-
-        # Then only for one of the sa emails
-        context = unittest.mock.Mock()
-        sa_filter = sa_emails[0]
-        response = self.grpc.ListHmacKeys(
-            storage_pb2.ListHmacKeysRequest(
-                project="projects/test-project-id", service_account_email=sa_filter
-            ),
-            context,
-        )
-        expected_access_ids = {
-            k.access_id for k in expected if k.service_account_email == sa_filter
-        }
-        self.assertEqual({k.access_id for k in response.hmac_keys}, expected_access_ids)
-
-        # Include deleted accounts (there are none, but should work)
-        context = unittest.mock.Mock()
-        sa_filter = sa_emails[0]
-        response = self.grpc.ListHmacKeys(
-            storage_pb2.ListHmacKeysRequest(
-                project="projects/test-project-id",
-                service_account_email=sa_filter,
-                show_deleted_keys=True,
-            ),
-            context,
-        )
-        expected_access_ids = {
-            k.access_id for k in expected if k.service_account_email == sa_filter
-        }
-        self.assertEqual({k.access_id for k in response.hmac_keys}, expected_access_ids)
-
-        # Missing or malformed project id is an error
-        context = unittest.mock.Mock()
-        _ = self.grpc.ListHmacKeys(storage_pb2.ListHmacKeysRequest(), context)
-        context.abort.assert_called_once_with(
-            grpc.StatusCode.INVALID_ARGUMENT, unittest.mock.ANY
-        )
-
-    def test_update_hmac_key(self):
-        sa_email = "test-sa@test-project-id.iam.gserviceaccount.com"
-        context = unittest.mock.Mock()
-        create = self.grpc.CreateHmacKey(
-            storage_pb2.CreateHmacKeyRequest(
-                project="projects/test-project-id",
-                service_account_email=sa_email,
-            ),
-            context,
-        )
-
-        metadata = storage_pb2.HmacKeyMetadata()
-        metadata.CopyFrom(create.metadata)
-        metadata.state = "INACTIVE"
-
-        context = unittest.mock.Mock()
-        response = self.grpc.UpdateHmacKey(
-            storage_pb2.UpdateHmacKeyRequest(
-                hmac_key=metadata, update_mask=field_mask_pb2.FieldMask(paths=["state"])
-            ),
-            context,
-        )
-        self.assertEqual(response.id, metadata.id)
-        self.assertEqual(response.access_id, metadata.access_id)
-        self.assertEqual(response.project, metadata.project)
-        self.assertEqual(response.service_account_email, metadata.service_account_email)
-        self.assertEqual(response.state, metadata.state)
-        self.assertEqual(response.create_time, metadata.create_time)
-
-        # Verify a missing, empty, or invalid update mask returns an error
-        for mask in [[], ["not-state"], ["state", "and-more"]]:
-            context = unittest.mock.Mock()
-            _ = self.grpc.UpdateHmacKey(
-                storage_pb2.UpdateHmacKeyRequest(
-                    hmac_key=metadata, update_mask=field_mask_pb2.FieldMask(paths=mask)
-                ),
-                context,
-            )
-            context.abort.assert_called_once_with(
-                grpc.StatusCode.INVALID_ARGUMENT, unittest.mock.ANY
-            )
-
-        # An empty metadata attribute is also an error
-        context = unittest.mock.Mock()
-        _ = self.grpc.UpdateHmacKey(
-            storage_pb2.UpdateHmacKeyRequest(
-                hmac_key=storage_pb2.HmacKeyMetadata(),
-                update_mask=field_mask_pb2.FieldMask(paths=["state"]),
-            ),
-            context,
-        )
-        context.abort.assert_called_once_with(
-            grpc.StatusCode.INVALID_ARGUMENT, unittest.mock.ANY
-        )
-
-    def test_update_hmac_key_bad_etag(self):
-        sa_email = "test-sa@test-project-id.iam.gserviceaccount.com"
-        context = unittest.mock.Mock()
-        create = self.grpc.CreateHmacKey(
-            storage_pb2.CreateHmacKeyRequest(
-                project="projects/test-project-id",
-                service_account_email=sa_email,
-            ),
-            context,
-        )
-
-        metadata = storage_pb2.HmacKeyMetadata()
-        metadata.CopyFrom(create.metadata)
-        metadata.state = "INACTIVE"
-        metadata.etag = "test-only-invalid"
-
-        context = unittest.mock.Mock()
-        response = self.grpc.UpdateHmacKey(
-            storage_pb2.UpdateHmacKeyRequest(
-                hmac_key=metadata, update_mask=field_mask_pb2.FieldMask(paths=["state"])
-            ),
-            context,
-        )
-        context.abort.assert_called_once_with(
-            grpc.StatusCode.FAILED_PRECONDITION, unittest.mock.ANY
-        )
-
     def test_run(self):
         port, server = testbench.grpc_server.run(0, self.db)
         self.assertNotEqual(port, 0)
diff --git a/tests/test_testbench_retry.py b/tests/test_testbench_retry.py
index a5b5f8a..0030993 100644
--- a/tests/test_testbench_retry.py
+++ b/tests/test_testbench_retry.py
@@ -30,6 +30,7 @@ import gcs
 import testbench
 from google.storage.v2 import storage_pb2
 from testbench import rest_server
+from tests.format_multipart_upload import format_multipart_upload
 
 UPLOAD_QUANTUM = 256 * 1024
 
@@ -655,6 +656,230 @@ class TestTestbenchRetry(unittest.TestCase):
         self.assertIn("size", create_rest)
         self.assertEqual(int(create_rest.get("size")), 2 * UPLOAD_QUANTUM)
 
+    def test_write_retry_test_stall_after_bytes(self):
+        # Create a new bucket
+        response = self.client.post(
+            "/storage/v1/b", data=json.dumps({"name": "bucket-name"})
+        )
+        self.assertEqual(response.status_code, 200)
+
+        # Setup a stall for reading back the object.
+        response = self.client.post(
+            "/retry_test",
+            data=json.dumps(
+                {
+                    "instructions": {
+                        "storage.objects.insert": [
+                            "stall-for-1s-after-250K",
+                            "stall-for-1s-after-300K",
+                        ]
+                    }
+                }
+            ),
+            content_type="application/json",
+        )
+        self.assertEqual(response.status_code, 200)
+        self.assertTrue(
+            response.headers.get("content-type").startswith("application/json")
+        )
+
+        create_rest = json.loads(response.data)
+        self.assertIn("id", create_rest)
+        test_id = create_rest.get("id")
+
+        # Initiate resumable upload
+        response = self.client.post(
+            "/upload/storage/v1/b/bucket-name/o",
+            query_string={"uploadType": "resumable", "name": "stall"},
+            content_type="application/json",
+        )
+        self.assertEqual(response.status_code, 200)
+
+        location = response.headers.get("location")
+        self.assertIn("upload_id=", location)
+        match = re.search(r"[&?]upload_id=([^&]+)", location)
+        self.assertIsNotNone(match, msg=location)
+        upload_id = match.group(1)
+
+        # Upload the first 256KiB chunk of data and trigger the stall.
+        chunk = self._create_block(UPLOAD_QUANTUM)
+        self.assertEqual(len(chunk), UPLOAD_QUANTUM)
+
+        start_time = time.perf_counter()
+        response = self.client.put(
+            f"/upload/storage/v1/b/bucket-name/o",
+            query_string={"upload_id": upload_id},
+            headers={
+                "content-range": "bytes 0-{len:d}/{obj_size:d}".format(
+                    len=UPLOAD_QUANTUM - 1, obj_size=2 * UPLOAD_QUANTUM
+                ),
+                "x-retry-test-id": test_id,
+            },
+            data=chunk,
+        )
+        end_time = time.perf_counter()
+        elapsed_time = end_time - start_time
+        self.assertGreater(elapsed_time, 1)
+        self.assertEqual(response.status_code, 308)
+
+        # Upload the second 256KiB chunk of data and trigger the stall again.
+        start_time = time.perf_counter()
+        chunk = self._create_block(UPLOAD_QUANTUM)
+        self.assertEqual(len(chunk), UPLOAD_QUANTUM)
+        response = self.client.put(
+            "/upload/storage/v1/b/bucket-name/o",
+            query_string={"upload_id": upload_id},
+            headers={
+                "content-range": "bytes 0-{len:d}/{obj_size:d}".format(
+                    len=2 * UPLOAD_QUANTUM - 1, obj_size=2 * UPLOAD_QUANTUM
+                ),
+                "x-retry-test-id": test_id,
+            },
+            data=chunk,
+        )
+        end_time = time.perf_counter()
+        elapsed_time = end_time - start_time
+        self.assertGreater(elapsed_time, 1)
+        self.assertEqual(response.status_code, 200, msg=response.data)
+
+        # Upload the second 256KiB chunk of data and check that stall not happen
+        start_time = time.perf_counter()
+        chunk = self._create_block(UPLOAD_QUANTUM)
+        self.assertEqual(len(chunk), UPLOAD_QUANTUM)
+        response = self.client.put(
+            "/upload/storage/v1/b/bucket-name/o",
+            query_string={"upload_id": upload_id},
+            headers={
+                "content-range": "bytes 0-{len:d}/{obj_size:d}".format(
+                    len=2 * UPLOAD_QUANTUM - 1, obj_size=2 * UPLOAD_QUANTUM
+                ),
+                "x-retry-test-id": test_id,
+            },
+            data=chunk,
+        )
+        end_time = time.perf_counter()
+        elapsed_time = end_time - start_time
+        self.assertLess(elapsed_time, 1)
+        self.assertEqual(response.status_code, 200, msg=response.data)
+
+    def test_write_retry_test_stall_single_shot(self):
+        # Create a new bucket
+        response = self.client.post(
+            "/storage/v1/b", data=json.dumps({"name": "bucket-name"})
+        )
+        self.assertEqual(response.status_code, 200)
+
+        # Setup a stall for reading back the object.
+        response = self.client.post(
+            "/retry_test",
+            data=json.dumps(
+                {
+                    "instructions": {
+                        "storage.objects.insert": [
+                            "stall-for-1s-after-250K",
+                        ]
+                    }
+                }
+            ),
+            content_type="application/json",
+        )
+        self.assertEqual(response.status_code, 200)
+        self.assertTrue(
+            response.headers.get("content-type").startswith("application/json")
+        )
+
+        create_rest = json.loads(response.data)
+        self.assertIn("id", create_rest)
+        test_id = create_rest.get("id")
+
+        # Upload the 256KiB of data and trigger the stall.
+        data = self._create_block(UPLOAD_QUANTUM)
+        self.assertEqual(len(data), UPLOAD_QUANTUM)
+
+        start_time = time.perf_counter()
+        boundary, payload = format_multipart_upload({}, data)
+        response = self.client.post(
+            "/upload/storage/v1/b/bucket-name/o",
+            query_string={"uploadType": "multipart", "name": "stall"},
+            content_type="multipart/related; boundary=" + boundary,
+            headers={
+                "x-retry-test-id": test_id,
+            },
+            data=payload,
+        )
+        end_time = time.perf_counter()
+        elapsed_time = end_time - start_time
+        self.assertEqual(response.status_code, 200)
+        self.assertGreater(elapsed_time, 1)
+
+        # Upload the data again and check that stall not happen.
+        start_time = time.perf_counter()
+        response = self.client.post(
+            "/upload/storage/v1/b/bucket-name/o",
+            query_string={"uploadType": "multipart", "name": "stall"},
+            content_type="multipart/related; boundary=" + boundary,
+            headers={
+                "x-retry-test-id": test_id,
+            },
+            data=payload,
+        )
+        end_time = time.perf_counter()
+        elapsed_time = end_time - start_time
+        self.assertLess(elapsed_time, 1)
+        self.assertEqual(response.status_code, 200)
+
+    def test_write_retry_test_stall_single_shot_while_upload_size_less_than_stall_size(
+        self,
+    ):
+        # Create a new bucket
+        response = self.client.post(
+            "/storage/v1/b", data=json.dumps({"name": "bucket-name"})
+        )
+        self.assertEqual(response.status_code, 200)
+
+        # Setup a stall for reading back the object.
+        response = self.client.post(
+            "/retry_test",
+            data=json.dumps(
+                {
+                    "instructions": {
+                        "storage.objects.insert": [
+                            "stall-for-1s-after-250K",
+                        ]
+                    }
+                }
+            ),
+            content_type="application/json",
+        )
+        self.assertEqual(response.status_code, 200)
+        self.assertTrue(
+            response.headers.get("content-type").startswith("application/json")
+        )
+
+        create_rest = json.loads(response.data)
+        self.assertIn("id", create_rest)
+        test_id = create_rest.get("id")
+
+        # Upload the 200KiB of data and check stall not happen.
+        data = self._create_block(200 * 1024)
+        self.assertEqual(len(data), 200 * 1024)
+
+        start_time = time.perf_counter()
+        boundary, payload = format_multipart_upload({}, data)
+        response = self.client.post(
+            "/upload/storage/v1/b/bucket-name/o",
+            query_string={"uploadType": "multipart", "name": "stall"},
+            content_type="multipart/related; boundary=" + boundary,
+            headers={
+                "x-retry-test-id": test_id,
+            },
+            data=payload,
+        )
+        end_time = time.perf_counter()
+        elapsed_time = end_time - start_time
+        self.assertEqual(response.status_code, 200)
+        self.assertLess(elapsed_time, 1)
+
 
 class TestTestbenchRetryGrpc(unittest.TestCase):
     def setUp(self):