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)
