def test_socket_get_logger(publish_safe, get_by_name, ZMQPubHandler): ( "SocketManager().get_logger() should return a logger" " with an attached ZMQPubHandler" ) # Background: get_by_name returns None get_by_name.return_value = None # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager without registered sockets manager = SocketManager(zmq, context) # When I call disconnect logger = manager.get_logger("foobar") # Then it should return false logger.should.be.a("logging.Logger") logger.handlers.should.equal([ZMQPubHandler.return_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)