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'