only handle after-bytes testbench instructions

testbench/rest_server.py

@@ -1038,7 +1038,7 @@ def resumable_upload_chunk(bucket_name):
             after_bytes,
             test_id,
         ) = testbench.common.get_retry_uploads_error_after_bytes(db, request)
-        if instruction or error_code:
+        if error_code or instruction == "return-503-after-256K":
             if instruction == "return-503-after-256K":
                 error_code = 503
                 after_bytes = 262144

tests/test_testbench_object_upload.py

@@ -617,6 +617,69 @@ def test_upload_pre_conditions_failure(self):
             )
             self.assertEqual(response.status_code, 412, msg=name)
 
+    def test_upload_resumable_w_fault_injection(self):
+        # Test fault injection "return-503-after-256K"
+        media = self._create_valid_chunk()
+        response = self.client.post(
+            "/upload/storage/v1/b/bucket-name/o",
+            query_string={"uploadType": "resumable", "name": "fox"},
+            content_type="application/json",
+            headers={
+                "x-goog-testbench-instructions": "return-503-after-256K",
+            },
+        )
+        self.assertEqual(response.status_code, 200)
+        location = response.headers.get("location")
+        self.assertIn("upload_id=", location)
+        match = re.search("[&?]upload_id=([^&]+)", location)
+        self.assertIsNotNone(match, msg=location)
+        upload_id = match.group(1)
+
+        response = self.client.put(
+            "/upload/storage/v1/b/bucket-name/o",
+            query_string={"upload_id": upload_id},
+            data=media,
+            headers={
+                "content-range": "bytes 0-{last:d}/{object_size:d}".format(
+                    last=UPLOAD_QUANTUM - 1, object_size=UPLOAD_QUANTUM
+                ),
+                "x-goog-testbench-instructions": "return-503-after-256K",
+            },
+        )
+        self.assertEqual(response.status_code, 503)
+
+        # Test fault injection "inject-upload-data-error"
+        response = self.client.post(
+            "/upload/storage/v1/b/bucket-name/o",
+            query_string={"uploadType": "resumable", "name": "zebra"},
+            content_type="application/json",
+            headers={
+                "x-goog-testbench-instructions": "inject-upload-data-error",
+            },
+        )
+        self.assertEqual(response.status_code, 200)
+        location = response.headers.get("location")
+        self.assertIn("upload_id=", location)
+        match = re.search("[&?]upload_id=([^&]+)", location)
+        self.assertIsNotNone(match, msg=location)
+        upload_id = match.group(1)
+
+        response = self.client.put(
+            "/upload/storage/v1/b/bucket-name/o",
+            query_string={"upload_id": upload_id},
+            data=media,
+            headers={
+                "x-goog-testbench-instructions": "inject-upload-data-error",
+            },
+        )
+        self.assertEqual(response.status_code, 200)
+
+        response = self.client.get(
+            "/download/storage/v1/b/bucket-name/o/zebra", query_string={"alt": "media"}
+        )
+        self.assertEqual(response.status_code, 200)
+        self.assertNotEqual(response.data.decode("utf-8"), media)
+
 
 if __name__ == "__main__":
     unittest.main()