def _grpc_data_channel_test(self, time_based_flush=False):
if time_based_flush:
data_servicer = data_plane.BeamFnDataServicer(
data_buffer_time_limit_ms=100)
else:
data_servicer = data_plane.BeamFnDataServicer()
worker_id = 'worker_0'
data_channel_service = \
data_servicer.get_conn_by_worker_id(worker_id)
server = grpc.server(UnboundedThreadPoolExecutor())
beam_fn_api_pb2_grpc.add_BeamFnDataServicer_to_server(
data_servicer, server)
test_port = server.add_insecure_port('[::]:0')
server.start()
grpc_channel = grpc.insecure_channel('localhost:%s' % test_port)
# Add workerId to the grpc channel
grpc_channel = grpc.intercept_channel(
grpc_channel, WorkerIdInterceptor(worker_id))
data_channel_stub = beam_fn_api_pb2_grpc.BeamFnDataStub(grpc_channel)
if time_based_flush:
data_channel_client = data_plane.GrpcClientDataChannel(
data_channel_stub, data_buffer_time_limit_ms=100)
else:
data_channel_client = data_plane.GrpcClientDataChannel(data_channel_stub)
try:
self._data_channel_test(
data_channel_service, data_channel_client, time_based_flush)
finally:
data_channel_client.close()
data_channel_service.close()
data_channel_client.wait()
data_channel_service.wait()
def test_grpc_data_channel(self):
data_servicer = data_plane.BeamFnDataServicer()
worker_id = 'worker_0'
data_channel_service = \
data_servicer.get_conn_by_worker_id(worker_id)
server = grpc.server(futures.ThreadPoolExecutor(max_workers=2))
beam_fn_api_pb2_grpc.add_BeamFnDataServicer_to_server(
data_servicer, server)
test_port = server.add_insecure_port('[::]:0')
server.start()
grpc_channel = grpc.insecure_channel('localhost:%s' % test_port)
# Add workerId to the grpc channel
grpc_channel = grpc.intercept_channel(grpc_channel,
WorkerIdInterceptor(worker_id))
data_channel_stub = beam_fn_api_pb2_grpc.BeamFnDataStub(grpc_channel)
data_channel_client = data_plane.GrpcClientDataChannel(
data_channel_stub)
try:
self._data_channel_test(data_channel_service, data_channel_client)
finally:
data_channel_client.close()
data_channel_service.close()
data_channel_client.wait()
data_channel_service.wait()
def test_grpc_data_channel(self):
data_channel_service = data_plane.GrpcServerDataChannel()
server = grpc.server(futures.ThreadPoolExecutor(max_workers=2))
beam_fn_api_pb2_grpc.add_BeamFnDataServicer_to_server(
data_channel_service, server)
test_port = server.add_insecure_port('[::]:0')
server.start()
data_channel_stub = beam_fn_api_pb2_grpc.BeamFnDataStub(
grpc.insecure_channel('localhost:%s' % test_port))
data_channel_client = data_plane.GrpcClientDataChannel(data_channel_stub)
try:
self._data_channel_test(data_channel_service, data_channel_client)
finally:
data_channel_client.close()
data_channel_service.close()
data_channel_client.wait()
data_channel_service.wait()