Make gapic's wrapper accept generic metadata (#4339)
diff --git a/google/api_core/gapic_v1/method.py b/google/api_core/gapic_v1/method.py
index e9a0bfa..0e84a84 100644
--- a/google/api_core/gapic_v1/method.py
+++ b/google/api_core/gapic_v1/method.py
@@ -94,16 +94,17 @@
timeout (google.api_core.timeout.Timeout): The default timeout
for the callable. If ``None``, this callable will not specify
a timeout argument to the low-level RPC method by default.
- user_agent_metadata (Tuple[str, str]): The user agent metadata key and
- value to provide to the RPC method. If ``None``, no additional
- metadata will be passed to the RPC method.
+ metadata (Sequence[Tuple[str, str]]): Additional metadata that is
+ provided to the RPC method on every invocation. This is merged with
+ any metadata specified during invocation. If ``None``, no
+ additional metadata will be passed to the RPC method.
"""
- def __init__(self, target, retry, timeout, user_agent_metadata=None):
+ def __init__(self, target, retry, timeout, metadata=None):
self._target = target
self._retry = retry
self._timeout = timeout
- self._user_agent_metadata = user_agent_metadata
+ self._metadata = metadata
def __call__(self, *args, **kwargs):
"""Invoke the low-level RPC with retry, timeout, and metadata."""
@@ -126,9 +127,9 @@
wrapped_func = _apply_decorators(self._target, [retry, timeout_])
# Add the user agent metadata to the call.
- if self._user_agent_metadata is not None:
- metadata = kwargs.get('metadata', [])
- metadata.append(self._user_agent_metadata)
+ if self._metadata is not None:
+ metadata = list(kwargs.get('metadata', []))
+ metadata.extend(self._metadata)
kwargs['metadata'] = metadata
return wrapped_func(*args, **kwargs)
@@ -219,11 +220,11 @@
func = grpc_helpers.wrap_errors(func)
if client_info is not None:
- user_agent_metadata = client_info.to_grpc_metadata()
+ user_agent_metadata = [client_info.to_grpc_metadata()]
else:
user_agent_metadata = None
return general_helpers.wraps(func)(
_GapicCallable(
func, default_retry, default_timeout,
- user_agent_metadata=user_agent_metadata))
+ metadata=user_agent_metadata))
diff --git a/tests/unit/gapic/test_method.py b/tests/unit/gapic/test_method.py
index 4034633..4592cb1 100644
--- a/tests/unit/gapic/test_method.py
+++ b/tests/unit/gapic/test_method.py
@@ -79,6 +79,21 @@
assert client_info.to_grpc_metadata() in metadata
+def test_invoke_wrapped_method_with_metadata():
+ method = mock.Mock(spec=['__call__'])
+
+ wrapped_method = google.api_core.gapic_v1.method.wrap_method(method)
+
+ wrapped_method(mock.sentinel.request, metadata=[('a', 'b')])
+
+ method.assert_called_once_with(mock.sentinel.request, metadata=mock.ANY)
+ metadata = method.call_args[1]['metadata']
+ # Metadata should have two items: the client info metadata and our custom
+ # metadata.
+ assert len(metadata) == 2
+ assert ('a', 'b') in metadata
+
+
@mock.patch('time.sleep')
def test_wrap_method_with_default_retry_and_timeout(unusued_sleep):
method = mock.Mock(