update handler retrieving retry test header

testbench/common.py

@@ -24,7 +24,6 @@
 import socket
 import struct
 import types
-from collections.abc import Iterable
 from functools import wraps
 
 import flask
@@ -742,12 +741,9 @@ def grpc_handle_retry_test_instruction(database, request, context, method):
     test_id = None
     if context is not None:
         if hasattr(context, "invocation_metadata") and isinstance(
-            context.invocation_metadata(), Iterable  # Handle mocks in tests
+            context.invocation_metadata(), tuple  # Handle mocks in tests
         ):
-            metadata = context.invocation_metadata()
-            if not isinstance(metadata, dict):
-                metadata = dict(metadata)
-            for key, value in metadata.items():
+            for key, value in context.invocation_metadata():
                 if key == "x-retry-test-id":
                     test_id = value
     # Validate retry instructions, method and request transport.

tests/test_testbench_retry.py

@@ -594,7 +594,7 @@ def test_grpc_retry_return_error(self):
 
         context = unittest.mock.Mock()
         context.invocation_metadata = unittest.mock.Mock(
-            return_value={"x-retry-test-id": create_rest.get("id")}
+            return_value=(("x-retry-test-id", create_rest.get("id")),)
         )
         response = self.grpc.GetBucket(
             storage_pb2.GetBucketRequest(name="projects/_/buckets/bucket-name"), context
@@ -620,7 +620,7 @@ def test_grpc_retry_reset_connection(self):
 
         context = unittest.mock.Mock()
         context.invocation_metadata = unittest.mock.Mock(
-            return_value={"x-retry-test-id": create_rest.get("id")}
+            return_value=(("x-retry-test-id", create_rest.get("id")),)
         )
         response = self.grpc.GetBucket(
             storage_pb2.GetBucketRequest(name="projects/_/buckets/bucket-name"), context