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_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)