Repository: googleapis/storage-testbench

Issue:
return error codes are not working for single shot uploads
Thanks for stopping by to let us know something could be better!

I was trying to run the following retry_test:
```
curl -X POST "http://localhost:5000/retry_test" -H 'Content-Type: application/json' -d '{"instructions":{"storage.objects.insert": ["return-503-after-3072K"]},"resources": ["BUCKET", "OBJECT"]}'
```
I observed that the error only occurred when uploading data larger than the chunk size (e.g., during resumable uploads). When uploading data smaller than the chunk size, the error did not happen.

Discussion:

Solution patch:
diff --git a/testbench/rest_server.py b/testbench/rest_server.py
index 72a0ae8..5e58296 100644
--- a/testbench/rest_server.py
+++ b/testbench/rest_server.py
@@ -979,17 +979,36 @@ 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)
+
+    # Handle errors for single-shot uploads.
+    testbench.common.extract_instruction(request, context=None)
+    (
+        error_code,
+        after_bytes,
+        test_id,
+    ) = testbench.common.get_retry_uploads_error_after_bytes(db, request)
+
+    if error_code and len(blob.media) >= after_bytes:
+        if test_id:
+            db.dequeue_next_instruction(test_id, "storage.objects.insert")
+        testbench.error.generic(
+            "Fault injected during a single-shot upload",
+            rest_code=error_code,
+            grpc_code=None,
+            context=None,
+        )
+
+    # Handle stall for single-shot 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,


Test patch:
diff --git a/tests/test_testbench_retry.py b/tests/test_testbench_retry.py
index 0030993..5a6fc4f 100644
--- a/tests/test_testbench_retry.py
+++ b/tests/test_testbench_retry.py
@@ -880,6 +880,112 @@ class TestTestbenchRetry(unittest.TestCase):
         self.assertEqual(response.status_code, 200)
         self.assertLess(elapsed_time, 1)
 
+    def test_retry_test_return_error_after_bytes_for_single_shot_upload(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": [
+                            "return-503-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)
+
+        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,
+        )
+        self.assertEqual(response.status_code, 503)
+
+        # Upload the data again and check that stall not happen.
+        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,
+        )
+        self.assertEqual(response.status_code, 200)
+
+    def test_write_retry_error_single_shot_while_upload_size_less_than_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 error for reading back the object.
+        response = self.client.post(
+            "/retry_test",
+            data=json.dumps(
+                {
+                    "instructions": {
+                        "storage.objects.insert": [
+                            "return-503-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 error not happen.
+        data = self._create_block(200 * 1024)
+        self.assertEqual(len(data), 200 * 1024)
+
+        boundary, payload = format_multipart_upload({}, data)
+        response = self.client.post(
+            "/upload/storage/v1/b/bucket-name/o",
+            query_string={"uploadType": "multipart", "name": "error"},
+            content_type="multipart/related; boundary=" + boundary,
+            headers={
+                "x-retry-test-id": test_id,
+            },
+            data=payload,
+        )
+        self.assertEqual(response.status_code, 200)
+
 
 class TestTestbenchRetryGrpc(unittest.TestCase):
     def setUp(self):