def test_execute_with_callback(self, mock_hook):
mocked_hook = mock.Mock()
callback = mock.Mock()
mock_hook.return_value = mocked_hook
mocked_hook.configure_mock(
**{'run.return_value': ["value1", "value2"]})
operator = GrpcOperator(stub_class=StubClass,
call_func="stream_call",
task_id="test_grpc",
response_callback=callback)
with mock.patch.object(operator.log, 'info') as mock_info:
operator.execute({})
mock_hook.assert_called_once_with("grpc_default",
interceptors=None,
custom_connection_func=None)
mocked_hook.run.assert_called_once_with(StubClass,
"stream_call",
data={},
streaming=False)
self.assertTrue(("'value1'",
"'value2'") not in mock_info.call_args_list)
mock_info.assert_any_call("Calling gRPC service")
callback.assert_any_call("value1", {})
callback.assert_any_call("value2", {})
def test_with_interceptors(self, mock_hook):
operator = GrpcOperator(
stub_class=StubClass, call_func="stream_call", interceptors=[], task_id="test_grpc",
)
operator.execute({})
mock_hook.assert_called_once_with("grpc_default", interceptors=[], custom_connection_func=None)
