def test_socket_close(get_by_name): ("SocketManager().close() should unregister, pop it " "from the list of connected sockets and then close") socket = get_by_name.return_value # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager manager = SocketManager(zmq, context) manager.poller = Mock(name='poller') manager.addresses = { 'foobar': 'wooot' } # When I call close manager.close("foobar") # Then it should have unregistered the socket with the poller manager.poller.unregister.assert_called_once_with(socket) # And the socket should have been removed from the internal addresses manager.addresses.should.be.empty # And socket.close should have been called socket.close.assert_called_once_with()
def test_socket_close(get_by_name): ("SocketManager().close() should unregister, pop it " "from the list of connected sockets and then close") socket = get_by_name.return_value # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager manager = SocketManager(zmq, context) manager.poller = Mock(name='poller') manager.addresses = {'foobar': 'wooot'} # When I call close manager.close("foobar") # Then it should have unregistered the socket with the poller manager.poller.unregister.assert_called_once_with(socket) # And the socket should have been removed from the internal addresses manager.addresses.should.be.empty # And socket.close should have been called socket.close.assert_called_once_with()
class GraphServer(object): def __init__(self, sleep_interval=3, pool_size=8): self.context = zmq.Context() self.sockets = SocketManager(zmq, self.context) self.sleep_interval = sleep_interval self.pool_size = pool_size self.pool = Pool(pool_size) self.allowed_to_run = Event() self.allowed_to_run.set() def should_run(self): return self.allowed_to_run.is_set() def listener_coroutine(self, socket_name, reply_bind_addr): # TODO: DEALER/ROUTER + internal multiplexing of request ids? self.sockets.ensure_and_bind(socket_name, zmq.REP, reply_bind_addr, zmq.POLLIN | zmq.POLLOUT) self.sockets.set_socket_option(socket_name, zmq.IDENTITY, socket_name) while self.should_run(): if not self.process_request(socket_name): logger.debug('waiting %s second%s', self.sleep_interval, self.sleep_interval == 1 and '' or 's') def process_request(self, socket_name): data = self.sockets.recv_safe(socket_name, timeout=self.sleep_interval) if data: logger.info('processed request') pprint(data) query = data['query'] result = schema.execute(query) if result.errors: payload = {'errors': "\n".join(map(unicode, result.errors))} else: payload = {'result': json.dumps(result.data, indent=2)} self.sockets.send_safe(socket_name, payload) return True return False def run(self, reply_bind_addr='tcp://*:6000'): logger.info('listening on %s', reply_bind_addr) self.pool.spawn(self.listener_coroutine, 'reply1', reply_bind_addr) self.pool.join(raise_error=True) def stop(self): self.allowed_to_run.clear() self.sockets.close('reply1') self.pool.kill()
def test_socket_close_no_socket(get_by_name): ("SocketManager().close() should do nothing if a socket " "with that given name does not exist") # Background: no sockets will be returned by get_by_name get_by_name.return_value = None # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager manager = SocketManager(zmq, context) manager.poller = Mock(name='poller') # When I call close manager.close("foobar") # Then it should not have called unregister manager.poller.unregister.called.should.be.false