fix(multipart): Make abort/complete return 200 on subsequent calls

testbench/rest_server.py

@@ -118,6 +118,8 @@ def xml_upload_part(bucket_name, object_name, upload_id, part_number):
     if getattr(upload, "kind", None) != "xml_multipart":
         testbench.error.notfound("Upload %s" % upload_id, None)
     _validate_xml_multipart_target(upload, bucket_name, object_name)
+    if getattr(upload, "aborted", False) or getattr(upload, "complete", False):
+        testbench.error.notfound("Upload %s" % upload_id, None)
     try:
         parsed_part_number = int(part_number)
     except (TypeError, ValueError):
@@ -141,6 +143,8 @@ def xml_list_parts(bucket_name, object_name, upload_id):
     if getattr(upload, "kind", None) != "xml_multipart":
         testbench.error.notfound("Upload %s" % upload_id, None)
     _validate_xml_multipart_target(upload, bucket_name, object_name)
+    if getattr(upload, "aborted", False) or getattr(upload, "complete", False):
+        testbench.error.notfound("Upload %s" % upload_id, None)
     max_parts = min(int(flask.request.args.get("max-parts", 1000)), 1000)
     marker = int(flask.request.args.get("part-number-marker", 0))
     sorted_nums = sorted(n for n in upload.parts if n > marker)
@@ -191,6 +195,22 @@ def xml_complete_multipart_upload(bucket_name, object_name, upload_id):
     if getattr(upload, "kind", None) != "xml_multipart":
         testbench.error.notfound("Upload %s" % upload_id, None)
     _validate_xml_multipart_target(upload, bucket_name, object_name)
+    if getattr(upload, "aborted", False):
+        testbench.error.notfound("Upload %s" % upload_id, None)
+    location = _xml_object_location(bucket_name, object_name)
+
+    # GCS returns 200 when complete is retried after a successful completion
+    # (e.g. response lost). Keep the session around marked complete so retries
+    # can reproduce the original success response instead of 404.
+    if getattr(upload, "complete", False):
+        multipart_etag = upload.multipart_etag
+        body = gcs_type.multipart_upload.build_complete_response_xml(
+            location, bucket_name, object_name, multipart_etag
+        )
+        response = flask.Response(body, status=200, content_type="application/xml")
+        response.headers["ETag"] = multipart_etag
+        return response
+
     requested = gcs_type.multipart_upload.parse_complete_request_xml(
         testbench.common.extract_media(flask.request)
     )
@@ -203,8 +223,9 @@ def xml_complete_multipart_upload(bucket_name, object_name, upload_id):
         context=None,
         preconditions=getattr(upload, "preconditions", []),
     )
-    db.delete_upload(upload_id, None)
-    location = _xml_object_location(bucket_name, object_name)
+    upload.complete = True
+    upload.multipart_etag = multipart_etag
+    upload.parts = {}
     body = gcs_type.multipart_upload.build_complete_response_xml(
         location, bucket_name, object_name, multipart_etag
     )
@@ -218,7 +239,16 @@ def xml_abort_multipart_upload(bucket_name, object_name, upload_id):
     if getattr(upload, "kind", None) != "xml_multipart":
         testbench.error.notfound("Upload %s" % upload_id, None)
     _validate_xml_multipart_target(upload, bucket_name, object_name)
-    db.delete_upload(upload_id, None)
+
+    # GCS returns 204 when abort is retried after a successful abort (e.g.
+    # response lost), and also when aborting a recently completed upload.
+    # Keep the session around marked aborted so retries stay successful
+    # instead of 404.
+    if getattr(upload, "aborted", False) or getattr(upload, "complete", False):
+        return flask.make_response("", 204)
+
+    upload.aborted = True
+    upload.parts = {}
     return flask.make_response("", 204)
 
 

tests/test_testbench_object_xml_multipart.py

@@ -200,6 +200,94 @@ def test_subdomain_happy_path(self):
     # ------------------------------------------------------------------
     # 3. Abort
     # ------------------------------------------------------------------
+    def test_complete_is_idempotent(self):
+        bucket = _create_bucket(self.client)
+        part = b"hello-complete-retry"
+        upload_id, _ = _initiate(self.client, bucket, "retry-complete.bin")
+        etag = _upload_part(
+            self.client, bucket, "retry-complete.bin", upload_id, 1, part
+        )
+
+        first = _complete(
+            self.client,
+            bucket,
+            "retry-complete.bin",
+            upload_id,
+            [(1, etag)],
+        )
+        self.assertEqual(first.status_code, 200, msg=first.data)
+        first_root = ET.fromstring(first.data)
+        first_etag = first_root.findtext(_ns("ETag"))
+        self.assertEqual(first.headers.get("ETag"), first_etag)
+
+        second = _complete(
+            self.client,
+            bucket,
+            "retry-complete.bin",
+            upload_id,
+            [(1, etag)],
+        )
+        self.assertEqual(second.status_code, 200, msg=second.data)
+        second_root = ET.fromstring(second.data)
+        self.assertEqual(second_root.findtext(_ns("ETag")), first_etag)
+        self.assertEqual(second.headers.get("ETag"), first_etag)
+
+        response = self.client.get("/%s/retry-complete.bin" % bucket)
+        self.assertEqual(response.status_code, 200, msg=response.data)
+        self.assertEqual(response.data, part)
+
+    def test_abort_is_idempotent(self):
+        bucket = _create_bucket(self.client)
+        upload_id, _ = _initiate(self.client, bucket, "retry-abort.txt")
+        _upload_part(self.client, bucket, "retry-abort.txt", upload_id, 1, b"data")
+
+        first = self.client.delete(
+            "/%s/retry-abort.txt" % bucket,
+            query_string={"uploadId": upload_id},
+        )
+        self.assertEqual(first.status_code, 204)
+
+        second = self.client.delete(
+            "/%s/retry-abort.txt" % bucket,
+            query_string={"uploadId": upload_id},
+        )
+        self.assertEqual(second.status_code, 204)
+
+        # Upload part after abort should still 404
+        response = self.client.put(
+            "/%s/retry-abort.txt" % bucket,
+            query_string={"uploadId": upload_id, "partNumber": "2"},
+            data=b"more",
+        )
+        self.assertEqual(response.status_code, 404)
+
+    def test_abort_after_complete_returns_204(self):
+        bucket = _create_bucket(self.client)
+        part = b"completed-then-aborted"
+        upload_id, _ = _initiate(self.client, bucket, "abort-after-complete.bin")
+        etag = _upload_part(
+            self.client, bucket, "abort-after-complete.bin", upload_id, 1, part
+        )
+        complete = _complete(
+            self.client,
+            bucket,
+            "abort-after-complete.bin",
+            upload_id,
+            [(1, etag)],
+        )
+        self.assertEqual(complete.status_code, 200, msg=complete.data)
+
+        response = self.client.delete(
+            "/%s/abort-after-complete.bin" % bucket,
+            query_string={"uploadId": upload_id},
+        )
+        self.assertEqual(response.status_code, 204)
+
+        # Object remains after aborting a completed upload.
+        response = self.client.get("/%s/abort-after-complete.bin" % bucket)
+        self.assertEqual(response.status_code, 200, msg=response.data)
+        self.assertEqual(response.data, part)
+
     def test_abort(self):
         bucket = _create_bucket(self.client)
         upload_id, _ = _initiate(self.client, bucket, "abort.txt")