Skip to content

Commit bc5a57c

Browse files
TulsishahTulsishah
authored
feat: Support error code for single shot upload (#699)
* support error for single shot upload * remove .idea files * lint fix * review comments * review comments * review comments * lint fix * lint fix * lint fix
1 parent 3ffa226 commit bc5a57c

File tree

2 files changed

+136
-11
lines changed

2 files changed

+136
-11
lines changed

‎testbench/rest_server.py‎

Lines changed: 30 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -979,17 +979,36 @@ def object_insert(bucket_name):
979979
blob, projection = gcs_type.object.Object.init_media(flask.request, bucket)
980980
elif upload_type == "multipart":
981981
blob, projection = gcs_type.object.Object.init_multipart(flask.request, bucket)
982-
# Handle stall for full uploads.
983-
testbench.common.extract_instruction(request, context=None)
984-
(
985-
stall_time,
986-
after_bytes,
987-
test_id,
988-
) = testbench.common.get_stall_uploads_after_bytes(db, request)
989-
if stall_time and len(blob.media) >= after_bytes:
990-
if test_id:
991-
db.dequeue_next_instruction(test_id, "storage.objects.insert")
992-
time.sleep(stall_time)
982+
983+
# Handle errors for single-shot uploads.
984+
testbench.common.extract_instruction(request, context=None)
985+
(
986+
error_code,
987+
after_bytes,
988+
test_id,
989+
) = testbench.common.get_retry_uploads_error_after_bytes(db, request)
990+
991+
if error_code and len(blob.media) >= after_bytes:
992+
if test_id:
993+
db.dequeue_next_instruction(test_id, "storage.objects.insert")
994+
testbench.error.generic(
995+
"Fault injected during a single-shot upload",
996+
rest_code=error_code,
997+
grpc_code=None,
998+
context=None,
999+
)
1000+
1001+
# Handle stall for single-shot uploads.
1002+
testbench.common.extract_instruction(request, context=None)
1003+
(
1004+
stall_time,
1005+
after_bytes,
1006+
test_id,
1007+
) = testbench.common.get_stall_uploads_after_bytes(db, request)
1008+
if stall_time and len(blob.media) >= after_bytes:
1009+
if test_id:
1010+
db.dequeue_next_instruction(test_id, "storage.objects.insert")
1011+
time.sleep(stall_time)
9931012

9941013
db.insert_object(
9951014
bucket_name,

‎tests/test_testbench_retry.py‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -880,6 +880,112 @@ def test_write_retry_test_stall_single_shot_while_upload_size_less_than_stall_si
880880
self.assertEqual(response.status_code, 200)
881881
self.assertLess(elapsed_time, 1)
882882

883+
def test_retry_test_return_error_after_bytes_for_single_shot_upload(self):
884+
# Create a new bucket
885+
response = self.client.post(
886+
"/storage/v1/b", data=json.dumps({"name": "bucket-name"})
887+
)
888+
self.assertEqual(response.status_code, 200)
889+
890+
# Setup a stall for reading back the object.
891+
response = self.client.post(
892+
"/retry_test",
893+
data=json.dumps(
894+
{
895+
"instructions": {
896+
"storage.objects.insert": [
897+
"return-503-after-250K",
898+
]
899+
}
900+
}
901+
),
902+
content_type="application/json",
903+
)
904+
self.assertEqual(response.status_code, 200)
905+
self.assertTrue(
906+
response.headers.get("content-type").startswith("application/json")
907+
)
908+
909+
create_rest = json.loads(response.data)
910+
self.assertIn("id", create_rest)
911+
test_id = create_rest.get("id")
912+
913+
# Upload the 256KiB of data and trigger the stall.
914+
data = self._create_block(UPLOAD_QUANTUM)
915+
self.assertEqual(len(data), UPLOAD_QUANTUM)
916+
917+
boundary, payload = format_multipart_upload({}, data)
918+
response = self.client.post(
919+
"/upload/storage/v1/b/bucket-name/o",
920+
query_string={"uploadType": "multipart", "name": "stall"},
921+
content_type="multipart/related; boundary=" + boundary,
922+
headers={
923+
"x-retry-test-id": test_id,
924+
},
925+
data=payload,
926+
)
927+
self.assertEqual(response.status_code, 503)
928+
929+
# Upload the data again and check that stall not happen.
930+
response = self.client.post(
931+
"/upload/storage/v1/b/bucket-name/o",
932+
query_string={"uploadType": "multipart", "name": "stall"},
933+
content_type="multipart/related; boundary=" + boundary,
934+
headers={
935+
"x-retry-test-id": test_id,
936+
},
937+
data=payload,
938+
)
939+
self.assertEqual(response.status_code, 200)
940+
941+
def test_write_retry_error_single_shot_while_upload_size_less_than_size(
942+
self,
943+
):
944+
# Create a new bucket
945+
response = self.client.post(
946+
"/storage/v1/b", data=json.dumps({"name": "bucket-name"})
947+
)
948+
self.assertEqual(response.status_code, 200)
949+
950+
# Setup a error for reading back the object.
951+
response = self.client.post(
952+
"/retry_test",
953+
data=json.dumps(
954+
{
955+
"instructions": {
956+
"storage.objects.insert": [
957+
"return-503-after-250K",
958+
]
959+
}
960+
}
961+
),
962+
content_type="application/json",
963+
)
964+
self.assertEqual(response.status_code, 200)
965+
self.assertTrue(
966+
response.headers.get("content-type").startswith("application/json")
967+
)
968+
969+
create_rest = json.loads(response.data)
970+
self.assertIn("id", create_rest)
971+
test_id = create_rest.get("id")
972+
973+
# Upload the 200KiB of data and check error not happen.
974+
data = self._create_block(200 * 1024)
975+
self.assertEqual(len(data), 200 * 1024)
976+
977+
boundary, payload = format_multipart_upload({}, data)
978+
response = self.client.post(
979+
"/upload/storage/v1/b/bucket-name/o",
980+
query_string={"uploadType": "multipart", "name": "error"},
981+
content_type="multipart/related; boundary=" + boundary,
982+
headers={
983+
"x-retry-test-id": test_id,
984+
},
985+
data=payload,
986+
)
987+
self.assertEqual(response.status_code, 200)
988+
883989

884990
class TestTestbenchRetryGrpc(unittest.TestCase):
885991
def setUp(self):

0 commit comments

Comments
 (0)