feat: add force_direct_connectivity behavior (#770)

* feat: add force_direct_connectivity behavior

* chore: add tests, readme and fix formatting

* style: fix formatting with black and isort

Author: cpriti-os

README.md

@@ -276,6 +276,7 @@ curl -H "x-retry-test-id: 1d05c20627844214a9ff7cbcf696317d" "http://localhost:91
 | stall-for-Ts-after-YK                     | [HTTP] Testbench will stall for T second after reading YKiB of downloaded/uploaded data, e.g. stall-for-10s-after-12K stalls after reading/writing 12KiB of data <br> [GRPC] Not supported
 | redirect-send-token-T                     | [HTTP] Unsupported [GRPC] Testbench will fail the RPC with `ABORTED` and include appropriate redirection error details.
 | redirect-send-handle-and-token-T          | [HTTP] Unsupported [GRPC] Testbench will fail the RPC with `ABORTED` and include appropriate redirection error details.
+| return-X-if-dp-enforced                   | [HTTP] Unsupported [GRPC] Testbench will fail with the equivalent gRPC error to the HTTP code provided for X, but only if DirectPath is enforced.
 
 ## Developing for the testbench
 

testbench/common.py

@@ -53,6 +53,7 @@
     r"redirect-send-handle-and-token-([a-z\-]+)$"
 )
 retry_expect_redirection_token = re.compile(r"redirect-expect-token-([a-z\-]+)$")
+retry_return_error_if_dp_enforced = re.compile(r"return-([0-9]+)-if-dp-enforced$")
 
 # A retry instruction to return a specific set of unreachable buckets.
 # For example, if the instruction is:
@@ -784,6 +785,23 @@ def grpc_handle_retry_test_instruction(database, request, context, method):
     ):
         return __get_default_response_fn
     next_instruction = database.peek_next_instruction(test_id, method)
+    enforced_match = retry_return_error_if_dp_enforced.match(next_instruction)
+    if enforced_match:
+        database.dequeue_next_instruction(test_id, method)
+        params = get_context_request_params(context)
+        if params and "force_direct_connectivity=ENFORCED" in params:
+            rest_code = enforced_match.group(1)
+            grpc_code = _grpc_forced_failure_from_http_instruction(rest_code)
+            msg = {
+                "error": {
+                    "message": "DirectPath is enforced but failed to establish connection: {}".format(
+                        grpc_code
+                    )
+                }
+            }
+            testbench.error.inject_error(context, rest_code, grpc_code, msg=msg)
+        return __get_default_response_fn
+
     error_code_matches = testbench.common.retry_return_error_code.match(
         next_instruction
     )

testbench/database.py

@@ -662,6 +662,7 @@ def __validate_injected_failure_description(self, failure):
             testbench.common.retry_return_handle_and_redirection_token,
             testbench.common.retry_expect_redirection_token,
             testbench.common.retry_return_unreachable_buckets,
+            testbench.common.retry_return_error_if_dp_enforced,
         ]:
             if expr.match(failure) is not None:
                 return

tests/test_testbench_retry.py

@@ -1442,6 +1442,56 @@ def test_grpc_bidiread_open_redirect(self):
         # Verify the early break occurred during first message only.
         self.assertEqual(len(responses), 0)
 
+    def test_grpc_retry_return_error_if_dp_enforced(self):
+        # Setup a retry test with 'return-503-if-dp-enforced'
+        response = self.rest_client.post(
+            "/retry_test",
+            data=json.dumps(
+                {
+                    "instructions": {
+                        "storage.buckets.get": ["return-503-if-dp-enforced"]
+                    },
+                    "transport": "GRPC",
+                },
+            ),
+        )
+        self.assertEqual(response.status_code, 200)
+        test_id = json.loads(response.data).get("id")
+
+        # Call WITH force_direct_connectivity=ENFORCED, expect 503 (UNAVAILABLE)
+        context = unittest.mock.Mock()
+        context.invocation_metadata.return_value = (
+            ("x-retry-test-id", test_id),
+            ("x-goog-request-params", "force_direct_connectivity=ENFORCED"),
+        )
+        self.grpc.GetBucket(
+            storage_pb2.GetBucketRequest(name="projects/_/buckets/bucket-name"), context
+        )
+        context.abort.assert_called_once_with(StatusCode.UNAVAILABLE, unittest.mock.ANY)
+
+        # Re-setup to test non-enforced behavior
+        response = self.rest_client.post(
+            "/retry_test",
+            data=json.dumps(
+                {
+                    "instructions": {
+                        "storage.buckets.get": ["return-503-if-dp-enforced"]
+                    },
+                    "transport": "GRPC",
+                }
+            ),
+        )
+        test_id = json.loads(response.data).get("id")
+
+        # Call WITHOUT enforcement, expect SUCCESS (instruction should still be consumed)
+        context = unittest.mock.Mock()
+        context.invocation_metadata.return_value = (("x-retry-test-id", test_id),)
+        self.grpc.GetBucket(
+            storage_pb2.GetBucketRequest(name="projects/_/buckets/bucket-name"), context
+        )
+        self.assertEqual(context.abort.call_count, 0)
+        self.assertIsNone(self.db.peek_next_instruction(test_id, "storage.buckets.get"))
+
     class _StatusAsCall:
         """_StatusAsCall wraps a status and pretends it is a client-side call"""