def test_socket_manager_publish_safe():
(
"SocketManager().publish_safe() should serialize "
"before sending, using the configured backend"
)
# Given a zmq mock
zmq = Mock()
# And a context
context = Mock()
# And a serializer
serializer = Mock(name="serializer")
serializer.pack.side_effect = lambda x: "<pac({})ked>".format(repr(x))
# And a socket manager
manager = SocketManager(zmq, context, serialization_backend=serializer)
# And a socket
socket = manager.create("foobar", zmq.REP)
# When I call .publish_safe()
manager.publish_safe("foobar", "topic", "PAYLOAD")
# Then it should have packed the payload before sending
serializer.pack.assert_called_once_with("PAYLOAD")
serializer.pack.return_value = "packed"
socket.send_multipart.assert_called_once_with(
[cast_bytes("topic"), cast_bytes("<pac('PAYLOAD')ked>")]
)
def test_socket_manager_publish_safe():
("SocketManager().publish_safe() should serialize "
"before sending, using the configured backend")
# Given a zmq mock
zmq = Mock()
# And a context
context = Mock()
# And a serializer
serializer = Mock(name='serializer')
# And a socket manager
manager = SocketManager(zmq, context, serialization_backend=serializer)
# And a socket
socket = manager.create('foobar', zmq.REP)
# When I call .publish_safe()
manager.publish_safe('foobar', 'topic', 'PAYLOAD')
# Then it should have packed the payload before sending
serializer.pack.assert_called_once_with('PAYLOAD')
packed = serializer.pack.return_value
socket.send_multipart.assert_called_once_with(['topic', packed])
def test_socket_manager_publish_safe():
("SocketManager().publish_safe() should serialize "
"before sending, using the configured backend")
# Given a zmq mock
zmq = Mock()
# And a context
context = Mock()
# And a serializer
serializer = Mock(name='serializer')
# And a socket manager
manager = SocketManager(zmq, context, serialization_backend=serializer)
# And a socket
socket = manager.create('foobar', zmq.REP)
# When I call .publish_safe()
manager.publish_safe('foobar', 'topic', 'PAYLOAD')
# Then it should have packed the payload before sending
serializer.pack.assert_called_once_with('PAYLOAD')
packed = serializer.pack.return_value
socket.send_multipart.assert_called_once_with(['topic', packed])
def test_socket_manager_publish_safe_not_ready():
("SocketManager.publish_safe should return False when the socket is not ready")
# Given a manager
manager = SocketManager(zmq, context)
# And a couple of sockets
manager.ensure_and_bind('foo', zmq.PUB, 'inproc://test.publisher.1', zmq.POLLOUT)
# When I call .publish_safe()
manager.publish_safe('foo', 'some-topic', {'some': 'value'})
class Step(object):
def __init__(self, pull_bind_address='tcp://127.0.0.1', subscriber_connect_address='tcp://127.0.0.1:6000', concurrency=100, timeout=1):
self.context = zmq.Context()
self.sockets = SocketManager(zmq, self.context)
self.sockets.create('pull-in', zmq.PULL)
# self.sockets.set_socket_option('pull-in', zmq.RCVHWM, concurrency)
self.sockets.create('events', zmq.PUB)
self.name = self.__class__.__name__
self.subscriber_connect_address = subscriber_connect_address
self._allowed_to_run = True
self.pool = gevent.pool.Pool(concurrency + 1)
self.timeout = timeout
self.pull_bind_address = pull_bind_address
self.id = str(uuid.uuid4())
self.logger = self.sockets.get_logger('events', 'logs', 'logs')
def listen(self):
_, self.address = self.sockets.bind_to_random_port('pull-in', zmq.POLLIN, local_address=self.pull_bind_address)
gevent.sleep(1)
def connect(self):
self.sockets.connect('events', self.subscriber_connect_address, zmq.POLLOUT)
gevent.sleep(1)
self.notify_available()
def notify_available(self):
self.send_event('available', self.to_dict())
def should_run(self):
gevent.sleep(0.1)
return self._allowed_to_run
def send_event(self, name, data):
self.sockets.publish_safe('events', name, data)
def dispatch(self, job):
try:
start = time.time()
job['job_started_at'] = start
self.send_event('job:started', job)
job['instructions'] = self.execute(job['instructions'])
job['job_finished_at'] = time.time()
self.logger.info("done processing %s", job)
except Exception as e:
job['job_finished_at'] = time.time()
job['error'] = serialized_exception(e)
self.send_event('job:error', job)
job['success'] = False
self.send_event('job:failed', job)
self.logger.exception('failed to execute job {id}'.format(**dict(job)))
finally:
end = time.time()
job['end'] = end
if job.get('instructions'):
self.send_event('job:success', job)
else:
self.send_event('job:failed', job)
def to_dict(self):
return dict(
id=self.id,
address=self.address,
job_type=self.job_type,
step_name=self.name,
)
def loop(self):
self.listen()
self.connect()
self.logger.info('listening for jobs on %s', self.address)
while self.should_run():
if self.pool.free_count() == 0:
self.logger.info('waiting for an execution slot')
self.pool.wait_available()
job = self.sockets.recv_safe('pull-in')
if job:
self.logger.info('received job')
self.pool.spawn(self.dispatch, job)
else:
self.notify_available()
gevent.sleep(1)