api_core: Add ChannelStub to grpc_helpers (#4705)

diff --git a/tests/unit/operations_v1/test_operations_client.py b/tests/unit/operations_v1/test_operations_client.py
index 1b6e6d9..69d4dfc 100644
--- a/tests/unit/operations_v1/test_operations_client.py
+++ b/tests/unit/operations_v1/test_operations_client.py
@@ -12,90 +12,64 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import mock
-
+from google.api_core import grpc_helpers
 from google.api_core import operations_v1
 from google.api_core import page_iterator
 from google.longrunning import operations_pb2
+from google.protobuf import empty_pb2
 
 
-def make_operations_stub(channel):
-    return mock.Mock(
-        spec=[
-            'GetOperation', 'DeleteOperation', 'ListOperations',
-            'CancelOperation'])
-
-
-operations_stub_patch = mock.patch(
-    'google.longrunning.operations_pb2.OperationsStub',
-    autospec=True,
-    side_effect=make_operations_stub)
-
-
-@operations_stub_patch
-def test_constructor(operations_stub):
-    stub = make_operations_stub(None)
-    operations_stub.side_effect = None
-    operations_stub.return_value = stub
-
-    client = operations_v1.OperationsClient(mock.sentinel.channel)
-
-    assert client.operations_stub == stub
-    operations_stub.assert_called_once_with(mock.sentinel.channel)
-
-
-@operations_stub_patch
-def test_get_operation(operations_stub):
-    client = operations_v1.OperationsClient(mock.sentinel.channel)
-    client.operations_stub.GetOperation.return_value = mock.sentinel.operation
+def test_get_operation():
+    channel = grpc_helpers.ChannelStub()
+    client = operations_v1.OperationsClient(channel)
+    channel.GetOperation.response = operations_pb2.Operation(name='meep')
 
     response = client.get_operation('name')
 
-    request = client.operations_stub.GetOperation.call_args[0][0]
-    assert isinstance(request, operations_pb2.GetOperationRequest)
-    assert request.name == 'name'
-
-    assert response == mock.sentinel.operation
+    assert len(channel.GetOperation.requests) == 1
+    assert channel.GetOperation.requests[0].name == 'name'
+    assert response == channel.GetOperation.response
 
 
-@operations_stub_patch
-def test_list_operations(operations_stub):
-    client = operations_v1.OperationsClient(mock.sentinel.channel)
+def test_list_operations():
+    channel = grpc_helpers.ChannelStub()
+    client = operations_v1.OperationsClient(channel)
     operations = [
         operations_pb2.Operation(name='1'),
         operations_pb2.Operation(name='2')]
     list_response = operations_pb2.ListOperationsResponse(
         operations=operations)
-    client.operations_stub.ListOperations.return_value = list_response
+    channel.ListOperations.response = list_response
 
     response = client.list_operations('name', 'filter')
 
     assert isinstance(response, page_iterator.Iterator)
     assert list(response) == operations
 
-    request = client.operations_stub.ListOperations.call_args[0][0]
+    assert len(channel.ListOperations.requests) == 1
+    request = channel.ListOperations.requests[0]
     assert isinstance(request, operations_pb2.ListOperationsRequest)
     assert request.name == 'name'
     assert request.filter == 'filter'
 
 
-@operations_stub_patch
-def test_delete_operation(operations_stub):
-    client = operations_v1.OperationsClient(mock.sentinel.channel)
+def test_delete_operation():
+    channel = grpc_helpers.ChannelStub()
+    client = operations_v1.OperationsClient(channel)
+    channel.DeleteOperation.response = empty_pb2.Empty()
 
     client.delete_operation('name')
 
-    request = client.operations_stub.DeleteOperation.call_args[0][0]
-    assert isinstance(request, operations_pb2.DeleteOperationRequest)
-    assert request.name == 'name'
+    assert len(channel.DeleteOperation.requests) == 1
+    assert channel.DeleteOperation.requests[0].name == 'name'
 
 
-@operations_stub_patch
-def test_cancel_operation(operations_stub):
-    client = operations_v1.OperationsClient(mock.sentinel.channel)
+def test_cancel_operation():
+    channel = grpc_helpers.ChannelStub()
+    client = operations_v1.OperationsClient(channel)
+    channel.CancelOperation.response = empty_pb2.Empty()
 
     client.cancel_operation('name')
 
-    request = client.operations_stub.CancelOperation.call_args[0][0]
-    assert isinstance(request, operations_pb2.CancelOperationRequest)
-    assert request.name == 'name'
+    assert len(channel.CancelOperation.requests) == 1
+    assert channel.CancelOperation.requests[0].name == 'name'
diff --git a/tests/unit/test_grpc_helpers.py b/tests/unit/test_grpc_helpers.py
index 6ee4062..de093e5 100644
--- a/tests/unit/test_grpc_helpers.py
+++ b/tests/unit/test_grpc_helpers.py
@@ -19,6 +19,7 @@
 from google.api_core import exceptions
 from google.api_core import grpc_helpers
 import google.auth.credentials
+from google.longrunning import operations_pb2
 
 
 def test__patch_callable_name():
@@ -186,3 +187,147 @@
         scopes=scopes)
 
     credentials.with_scopes.assert_called_once_with(scopes)
+
+
+class TestChannelStub(object):
+
+    def test_single_response(self):
+        channel = grpc_helpers.ChannelStub()
+        stub = operations_pb2.OperationsStub(channel)
+        expected_request = operations_pb2.GetOperationRequest(name='meep')
+        expected_response = operations_pb2.Operation(name='moop')
+
+        channel.GetOperation.response = expected_response
+
+        response = stub.GetOperation(expected_request)
+
+        assert response == expected_response
+        assert channel.requests == [('GetOperation', expected_request)]
+        assert channel.GetOperation.requests == [expected_request]
+
+    def test_no_response(self):
+        channel = grpc_helpers.ChannelStub()
+        stub = operations_pb2.OperationsStub(channel)
+        expected_request = operations_pb2.GetOperationRequest(name='meep')
+
+        with pytest.raises(ValueError) as exc_info:
+            stub.GetOperation(expected_request)
+
+        assert exc_info.match('GetOperation')
+
+    def test_missing_method(self):
+        channel = grpc_helpers.ChannelStub()
+
+        with pytest.raises(AttributeError):
+            channel.DoesNotExist.response
+
+    def test_exception_response(self):
+        channel = grpc_helpers.ChannelStub()
+        stub = operations_pb2.OperationsStub(channel)
+        expected_request = operations_pb2.GetOperationRequest(name='meep')
+
+        channel.GetOperation.response = RuntimeError()
+
+        with pytest.raises(RuntimeError):
+            stub.GetOperation(expected_request)
+
+    def test_callable_response(self):
+        channel = grpc_helpers.ChannelStub()
+        stub = operations_pb2.OperationsStub(channel)
+        expected_request = operations_pb2.GetOperationRequest(name='meep')
+        expected_response = operations_pb2.Operation(name='moop')
+
+        on_get_operation = mock.Mock(
+            spec=('__call__',), return_value=expected_response)
+
+        channel.GetOperation.response = on_get_operation
+
+        response = stub.GetOperation(expected_request)
+
+        assert response == expected_response
+        on_get_operation.assert_called_once_with(expected_request)
+
+    def test_multiple_responses(self):
+        channel = grpc_helpers.ChannelStub()
+        stub = operations_pb2.OperationsStub(channel)
+        expected_request = operations_pb2.GetOperationRequest(name='meep')
+        expected_responses = [
+            operations_pb2.Operation(name='foo'),
+            operations_pb2.Operation(name='bar'),
+            operations_pb2.Operation(name='baz'),
+        ]
+
+        channel.GetOperation.responses = iter(expected_responses)
+
+        response1 = stub.GetOperation(expected_request)
+        response2 = stub.GetOperation(expected_request)
+        response3 = stub.GetOperation(expected_request)
+
+        assert response1 == expected_responses[0]
+        assert response2 == expected_responses[1]
+        assert response3 == expected_responses[2]
+        assert channel.requests == [('GetOperation', expected_request)] * 3
+        assert channel.GetOperation.requests == [expected_request] * 3
+
+        with pytest.raises(StopIteration):
+            stub.GetOperation(expected_request)
+
+    def test_multiple_responses_and_single_response_error(self):
+        channel = grpc_helpers.ChannelStub()
+        stub = operations_pb2.OperationsStub(channel)
+        channel.GetOperation.responses = []
+        channel.GetOperation.response = mock.sentinel.response
+
+        with pytest.raises(ValueError):
+            stub.GetOperation(operations_pb2.GetOperationRequest())
+
+    def test_call_info(self):
+        channel = grpc_helpers.ChannelStub()
+        stub = operations_pb2.OperationsStub(channel)
+        expected_request = operations_pb2.GetOperationRequest(name='meep')
+        expected_response = operations_pb2.Operation(name='moop')
+        expected_metadata = [('red', 'blue'), ('two', 'shoe')]
+        expected_credentials = mock.sentinel.credentials
+        channel.GetOperation.response = expected_response
+
+        response = stub.GetOperation(
+            expected_request, timeout=42, metadata=expected_metadata,
+            credentials=expected_credentials)
+
+        assert response == expected_response
+        assert channel.requests == [('GetOperation', expected_request)]
+        assert channel.GetOperation.calls == [
+            (expected_request, 42, expected_metadata, expected_credentials)]
+
+    def test_unary_unary(self):
+        channel = grpc_helpers.ChannelStub()
+        method_name = 'GetOperation'
+        callable_stub = channel.unary_unary(method_name)
+        assert callable_stub._method == method_name
+        assert callable_stub._channel == channel
+
+    def test_unary_stream(self):
+        channel = grpc_helpers.ChannelStub()
+        method_name = 'GetOperation'
+        callable_stub = channel.unary_stream(method_name)
+        assert callable_stub._method == method_name
+        assert callable_stub._channel == channel
+
+    def test_stream_unary(self):
+        channel = grpc_helpers.ChannelStub()
+        method_name = 'GetOperation'
+        callable_stub = channel.stream_unary(method_name)
+        assert callable_stub._method == method_name
+        assert callable_stub._channel == channel
+
+    def test_stream_stream(self):
+        channel = grpc_helpers.ChannelStub()
+        method_name = 'GetOperation'
+        callable_stub = channel.stream_stream(method_name)
+        assert callable_stub._method == method_name
+        assert callable_stub._channel == channel
+
+    def test_subscribe_unsubscribe(self):
+        channel = grpc_helpers.ChannelStub()
+        assert channel.subscribe(None) is None
+        assert channel.unsubscribe(None) is None