unit test fix

googleapis · Oct 27, 2024 · 328d6e1 · 328d6e1
1 parent 34b098a
commit 328d6e1
Showing 1 changed file with 29 additions and 24 deletions.
diff --git a/tests/test_testbench_retry.py b/tests/test_testbench_retry.py
@@ -828,7 +828,9 @@ def test_write_retry_test_stall_single_shot(self):
         self.assertLess(elapsed_time, 1)
         self.assertEqual(response.status_code, 200)
 
-    def test_retry_test_return_error_after_bytes_for_single_shot_upload(self):
+    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"})
@@ -842,7 +844,7 @@ def test_retry_test_return_error_after_bytes_for_single_shot_upload(self):
                 {
                     "instructions": {
                         "storage.objects.insert": [
-                            "return-503-after-250K",
+                            "stall-for-1s-after-250K",
                         ]
                     }
                 }
@@ -858,11 +860,12 @@ def test_retry_test_return_error_after_bytes_for_single_shot_upload(self):
         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)
+        # Upload the 200KiB of data and check stall not happen.
+        data = self._create_block(200 * 1024)
+        self.assertEqual(len(data), 200 * 1024)
 
         boundary, payload = format_multipart_upload({}, data)
+        start_time = time.perf_counter()
         response = self.client.post(
             "/upload/storage/v1/b/bucket-name/o",
             query_string={"uploadType": "multipart", "name": "stall"},
@@ -872,23 +875,13 @@ def test_retry_test_return_error_after_bytes_for_single_shot_upload(self):
             },
             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,
-        )
+        end_time = time.perf_counter()
+        elapsed_time = end_time - start_time
+        self.assertEqual(response.status_code, 200)
+        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,
-    ):
+    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"})
@@ -902,7 +895,7 @@ def test_write_retry_test_stall_single_shot_while_upload_size_less_than_stall_si
                 {
                     "instructions": {
                         "storage.objects.insert": [
-                            "stall-for-1s-after-250K",
+                            "return-503-after-250K",
                         ]
                     }
                 }
@@ -918,9 +911,9 @@ def test_write_retry_test_stall_single_shot_while_upload_size_less_than_stall_si
         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)
+        # 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(
@@ -932,6 +925,18 @@ def test_write_retry_test_stall_single_shot_while_upload_size_less_than_stall_si
             },
             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(