def test_grpc_data_runtime(mocker): args = set_pea_parser().parse_args([]) handle_mock = mocker.Mock() def start_runtime(args, handle_mock): with GRPCDataRuntime(args) as runtime: runtime._data_request_handler.handle = handle_mock runtime.run_forever() runtime_thread = Thread( target=start_runtime, args=( args, handle_mock, ), ) runtime_thread.start() assert GRPCDataRuntime.wait_for_ready_or_shutdown( timeout=5.0, ctrl_address=f'{args.host}:{args.port_in}', shutdown_event=Event()) Grpclet._create_grpc_stub(f'{args.host}:{args.port_in}', is_async=False).Call(_create_test_data_message()) time.sleep(0.1) handle_mock.assert_called() GRPCDataRuntime.cancel(f'{args.host}:{args.port_in}') runtime_thread.join() assert not GRPCDataRuntime.is_ready(f'{args.host}:{args.port_in}')
def test_grpc_data_runtime(mocker): args = set_pea_parser().parse_args([]) handle_mock = multiprocessing.Event() cancel_event = multiprocessing.Event() def start_runtime(args, handle_mock, cancel_event): with GRPCDataRuntime(args, cancel_event) as runtime: runtime._data_request_handler.handle = ( lambda *args, **kwargs: handle_mock.set() ) runtime.run_forever() runtime_thread = Process( target=start_runtime, args=(args, handle_mock, cancel_event), daemon=True, ) runtime_thread.start() assert GRPCDataRuntime.wait_for_ready_or_shutdown( timeout=5.0, ctrl_address=f'{args.host}:{args.port_in}', shutdown_event=Event() ) Grpclet._create_grpc_stub(f'{args.host}:{args.port_in}', is_async=False).Call( _create_test_data_message() ) time.sleep(0.1) assert handle_mock.is_set() GRPCDataRuntime.cancel(cancel_event) runtime_thread.join() assert not GRPCDataRuntime.is_ready(f'{args.host}:{args.port_in}')
def test_grpc_data_runtime_waits_for_pending_messages_shutdown(close_method): args = set_pea_parser().parse_args([]) cancel_event = multiprocessing.Event() handler_closed_event = multiprocessing.Event() slow_executor_block_time = 1.0 pending_requests = 3 sent_queue = multiprocessing.Queue() def start_runtime(args, cancel_event, sent_queue, handler_closed_event): with GRPCDataRuntime(args, cancel_event) as runtime: runtime._data_request_handler.handle = lambda *args, **kwargs: time.sleep( slow_executor_block_time ) runtime._data_request_handler.close = ( lambda *args, **kwargs: handler_closed_event.set() ) async def mock(msg): sent_queue.put('') runtime._grpclet.send_message = mock runtime.run_forever() runtime_thread = Process( target=start_runtime, args=(args, cancel_event, sent_queue, handler_closed_event), daemon=True, ) runtime_thread.start() assert GRPCDataRuntime.wait_for_ready_or_shutdown( timeout=5.0, ctrl_address=f'{args.host}:{args.port_in}', shutdown_event=Event() ) request_start_time = time.time() for i in range(pending_requests): Grpclet._create_grpc_stub(f'{args.host}:{args.port_in}', is_async=False).Call( _create_test_data_message() ) time.sleep(0.1) if close_method == 'TERMINATE': runtime_thread.terminate() else: GRPCDataRuntime.cancel(cancel_event) assert not handler_closed_event.is_set() runtime_thread.join() assert ( time.time() - request_start_time >= slow_executor_block_time * pending_requests ) assert sent_queue.qsize() == pending_requests assert handler_closed_event.is_set() assert not GRPCDataRuntime.is_ready(f'{args.host}:{args.port_in}')
def test_executor_runtimes(signal, tmpdir, grpc_data_requests): import time args = set_pea_parser().parse_args([]) def run(args): args.uses = { 'jtype': 'DummyExecutor', 'with': { 'dir': str(tmpdir) }, 'metas': { 'workspace': str(tmpdir) }, } args.grpc_data_requests = grpc_data_requests executor_native(args) process = multiprocessing.Process(target=run, args=(args, )) process.start() time.sleep(0.5) if grpc_data_requests: Grpclet._create_grpc_stub(f'{args.host}:{args.port_in}', is_async=False).Call( _create_test_data_message()) else: socket = zmq.Context().socket(zmq.PUSH) socket.connect(f'tcp://localhost:{args.port_in}') socket.send_multipart(_create_test_data_message().dump()) time.sleep(0.1) os.kill(process.pid, signal) process.join() with open(f'{tmpdir}/test.txt', 'r') as fp: output = fp.read() split = output.split(';') assert split[0] == 'proper close' assert split[1] == '1'