commit 2576a1b260b795651c50a07f4ddebabe7bef5bdd
author Christopher Wilcox <crwilcox@google.com> Wed Mar 27 10:58:39 2019 -0700
committer GitHub <noreply@github.com> Wed Mar 27 10:58:39 2019 -0700
tree 52ee497e4e5034791fc9681ed88b1db87aae58e7
parent 10a694ec945ada42c3ccecc8fe3877bd83ce26d8

Allow passing metadata as part of creating a bidi (#7514)

* allows providing rpc metadata for bidi streams

google/api_core/bidi.py

diff --git a/google/api_core/bidi.py b/google/api_core/bidi.py
index be24d32..053363a 100644
--- a/google/api_core/bidi.py
+++ b/google/api_core/bidi.py
@@ -147,7 +147,11 @@
 
         initial_request = example_pb2.StreamingRpcRequest(
             setting='example')
-        rpc = BidiRpc(stub.StreamingRpc, initial_request=initial_request)
+        rpc = BidiRpc(
+            stub.StreamingRpc,
+            initial_request=initial_request,
+            metadata=[('name', 'value')]
+        )
 
         rpc.open()
 
@@ -165,11 +169,14 @@
                 Callable[None, protobuf.Message]]): The initial request to
             yield. This is useful if an initial request is needed to start the
             stream.
+        metadata (Sequence[Tuple(str, str)]): RPC metadata to include in
+            the request.
     """
 
-    def __init__(self, start_rpc, initial_request=None):
+    def __init__(self, start_rpc, initial_request=None, metadata=None):
         self._start_rpc = start_rpc
         self._initial_request = initial_request
+        self._rpc_metadata = metadata
         self._request_queue = queue.Queue()
         self._request_generator = None
         self._is_active = False
@@ -200,7 +207,7 @@
         request_generator = _RequestQueueGenerator(
             self._request_queue, initial_request=self._initial_request
         )
-        call = self._start_rpc(iter(request_generator))
+        call = self._start_rpc(iter(request_generator), metadata=self._rpc_metadata)
 
         request_generator.call = call
 
@@ -288,10 +295,14 @@
         initial_request = example_pb2.StreamingRpcRequest(
             setting='example')
 
-        rpc = ResumeableBidiRpc(
+        metadata = [('header_name', 'value')]
+
+        rpc = ResumableBidiRpc(
             stub.StreamingRpc,
+            should_recover=should_recover,
             initial_request=initial_request,
-            should_recover=should_recover)
+            metadata=metadata
+        )
 
         rpc.open()
 
@@ -310,10 +321,12 @@
         should_recover (Callable[[Exception], bool]): A function that returns
             True if the stream should be recovered. This will be called
             whenever an error is encountered on the stream.
+        metadata Sequence[Tuple(str, str)]: RPC metadata to include in
+            the request.
     """
 
-    def __init__(self, start_rpc, should_recover, initial_request=None):
-        super(ResumableBidiRpc, self).__init__(start_rpc, initial_request)
+    def __init__(self, start_rpc, should_recover, initial_request=None, metadata=None):
+        super(ResumableBidiRpc, self).__init__(start_rpc, initial_request, metadata)
         self._should_recover = should_recover
         self._operational_lock = threading.RLock()
         self._finalized = False

noxfile.py

diff --git a/noxfile.py b/noxfile.py
index fbc51ee..06e78c1 100644
--- a/noxfile.py
+++ b/noxfile.py
@@ -88,7 +88,7 @@
 def pytype(session):
     """Run type-checking."""
     session.install(
-        ".", "grpcio >= 1.8.2", "grpcio-gcp >= 0.2.2", "pytype >= 2018.9.26"
+        ".", "grpcio >= 1.8.2", "grpcio-gcp >= 0.2.2", "pytype >= 2019.3.21"
     )
     session.run("pytype")
 

tests/unit/test_bidi.py

diff --git a/tests/unit/test_bidi.py b/tests/unit/test_bidi.py
index 8163699..08d2021 100644
--- a/tests/unit/test_bidi.py
+++ b/tests/unit/test_bidi.py
@@ -125,9 +125,10 @@
     call = mock.create_autospec(_CallAndFuture, instance=True)
     rpc = mock.create_autospec(grpc.StreamStreamMultiCallable, instance=True)
 
-    def rpc_side_effect(request):
+    def rpc_side_effect(request, metadata=None):
         call.is_active.return_value = True
         call.request = request
+        call.metadata = metadata
         return call
 
     rpc.side_effect = rpc_side_effect
@@ -172,6 +173,15 @@
 
         callback.assert_called_once_with(mock.sentinel.future)
 
+    def test_metadata(self):
+        rpc, call = make_rpc()
+        bidi_rpc = bidi.BidiRpc(rpc, metadata=mock.sentinel.A)
+        assert bidi_rpc._rpc_metadata == mock.sentinel.A
+
+        bidi_rpc.open()
+        assert bidi_rpc.call == call
+        assert bidi_rpc.call.metadata == mock.sentinel.A
+
     def test_open(self):
         rpc, call = make_rpc()
         bidi_rpc = bidi.BidiRpc(rpc)