unit test for metadata intialize

doc comment improvements
googleapis · Mar 26, 2019 · 4f16bdb · 4f16bdb
1 parent 56b57e3
commit 4f16bdb
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 10 deletions.
diff --git a/api_core/google/api_core/bidi.py b/api_core/google/api_core/bidi.py
@@ -147,7 +147,11 @@ class BidiRpc(object):
 
         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,14 +169,14 @@ class BidiRpc(object):
                 Callable[None, protobuf.Message]]): The initial request to
             yield. This is useful if an initial request is needed to start the
             stream.
-        rpc_metadata Sequence[Tuple(str, str)]: RPC metadata to include in
+        metadata (Sequence[Tuple(str, str)]): RPC metadata to include in
             the request.
     """
 
-    def __init__(self, start_rpc, initial_request=None, rpc_metadata=None):
+    def __init__(self, start_rpc, initial_request=None, metadata=None):
         self._start_rpc = start_rpc
         self._initial_request = initial_request
-        self._rpc_metadata = rpc_metadata
+        self._rpc_metadata = metadata
         self._request_queue = queue.Queue()
         self._request_generator = None
         self._is_active = False
@@ -291,10 +295,14 @@ def should_recover(exc):
         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()
 
@@ -313,13 +321,12 @@ def should_recover(exc):
         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.
-           def _rpc_metadata(self):
-        rpc_metadata Sequence[Tuple(str, str)]: RPC metadata to include in
+        metadata Sequence[Tuple(str, str)]: RPC metadata to include in
             the request.
     """
 
-    def __init__(self, start_rpc, should_recover, initial_request=None, rpc_metadata=None):
-        super(ResumableBidiRpc, self).__init__(start_rpc, initial_request, rpc_metadata)
+    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

diff --git a/api_core/tests/unit/test_bidi.py b/api_core/tests/unit/test_bidi.py
@@ -173,6 +173,15 @@ def test_done_callbacks(self):
 
         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)