def test_simple_run(self, mock_get_conn, mock_get_connection): conn = get_airflow_connection() mock_get_connection.return_value = conn mocked_channel = mock.Mock() mocked_channel.__enter__ = mock.Mock(return_value=(mock.Mock(), None)) mocked_channel.__exit__ = mock.Mock(return_value=None) hook = GrpcHook("grpc_default") mock_get_conn.return_value = mocked_channel response = hook.run(StubClass, "single_call", data={'data': 'hello'}) self.assertEqual(next(response), "hello")
def test_stream_run(self, mock_get_conn, mock_get_connection): conn = get_airflow_connection() mock_get_connection.return_value = conn mocked_channel = mock.Mock() mocked_channel.__enter__ = mock.Mock(return_value=(mock.Mock(), None)) mocked_channel.__exit__ = mock.Mock(return_value=None) hook = GrpcHook("grpc_default") mock_get_conn.return_value = mocked_channel response = hook.run(StubClass, "stream_call", data={'data': ['hello!', "hi"]}) assert next(response) == ["streaming", "call"]