class WorkerWithSocket: def __init__(self, attributes=("send_and_recv",)): self._enter_attributes = attributes self.worker = Worker("tcp://localhost", print_local=True) def __enter__(self): self.worker.start() ctx = zmq.Context() self.socket = WorkerConnection(ctx) self.socket.bind("tcp://*:23456") assert self.socket.poll(1000, zmq.POLLIN), "No response from started worker" self.worker_id, connect_message = self.socket.recv_response() assert connect_message == Commands.CONNECT, connect_message assert "worker-" in self.worker_id.decode("ascii") assert self.worker.is_alive() self.stdsocket = ctx.socket(zmq.PULL) self.stdsocket.bind("tcp://*:23457") attrs = [getattr(self, attr) for attr in self._enter_attributes] if len(attrs) == 1: return attrs[0] return attrs def __exit__(self, arg1, arg2, arg3): assert self.worker.is_alive(), "Process Died During Testing" self.socket.send_cmd(self.worker_id, Commands.DIE) self.worker.join() def send_and_recv(self, cmd, *args, timeout=1000, expect=None, **kwargs): self.socket.send_cmd(self.worker_id, cmd, *args) assert self.socket.poll(timeout, zmq.POLLIN), "Timeout waiting for response to: {}".format(args) worker_id, response = self.socket.recv_response(**kwargs) assert worker_id == self.worker_id, worker_id err_msg = "Invalid response from '{}' command: {}".format(Commands.items()[cmd], response) if expect is not None and callable(expect): assert expect(response), err_msg elif expect is not None: assert response == expect, err_msg return response
def __init__(self, attributes=("send_and_recv",)): self._enter_attributes = attributes self.worker = Worker("tcp://localhost", print_local=True)