def test_socket_manager_can_poll_asynchronously():
("SocketManager should leverage a non-blocking socket "
"collection can be used seamlessly in a blocking fashion [TCP SOCKET]")
# Given a socket manager for a server
server = SocketManager(zmq, context)
# And a reply socket listening on a tcp port
server.ensure_and_bind('reply-server', zmq.REP, 'tcp://0.0.0.0:3458', zmq.POLLIN | zmq.POLLOUT)
# And a socket manager for a client
client = SocketManager(zmq, context)
# And a request client connected to the server
client.ensure_and_connect('request-client', zmq.REQ, 'tcp://0.0.0.0:3458', zmq.POLLIN | zmq.POLLOUT)
# And send a request from the client
requester = client.wait_until_ready('request-client', zmq.POLLOUT, timeout=2)
requester.send_json({'client': 42})
# Then I should receive a request from the client
replier = server.wait_until_ready('reply-server', zmq.POLLIN, timeout=2)
replier.should_not.be.none
request = replier.recv_json()
# And the request should be the one that the client just sent
request.should.equal({'client': 42})
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 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_ensure_and_bind(get_or_create, bind, engage):
(
"SocketManager().ensure_and_bind() should ensure the socket existence then "
"bind to the given address."
)
# Given a zmq mock
zmq = Mock()
# And a context
context = Mock()
# And a socket manager
manager = SocketManager(zmq, context)
# When I call ensure_and_bind
result = manager.ensure_and_bind(
"foobar", "SOME_TYPE", "inproc://whatevs", "some-mechanism"
)
# Then it should have returned a socket from bind
result.should.equal(bind.return_value)
# And get_or_create should have been called correctly
get_or_create.assert_called_once_with(
"foobar", "SOME_TYPE", "some-mechanism"
)
# And bind should have been called with the given address
bind.assert_called_once_with("foobar", "inproc://whatevs", "some-mechanism")
# And engage should have been called with the default value
engage.assert_called_once_with()
def test_socket_ensure_and_bind(get_or_create, bind, engage):
("SocketManager().ensure_and_bind() should ensure the socket existence then "
"bind to the given address.")
# Given a zmq mock
zmq = Mock()
# And a context
context = Mock()
# And a socket manager
manager = SocketManager(zmq, context)
# When I call ensure_and_bind
result = manager.ensure_and_bind("foobar", "SOME_TYPE", 'inproc://whatevs',
'some-mechanism')
# Then it should have returned a socket from bind
result.should.equal(bind.return_value)
# And get_or_create should have been called correctly
get_or_create.assert_called_once_with('foobar', 'SOME_TYPE',
'some-mechanism')
# And bind should have been called with the given address
bind.assert_called_once_with('foobar', 'inproc://whatevs',
'some-mechanism')
# And engage should have been called with the default value
engage.assert_called_once_with()
def test_socket_manager_subscribe():
("SocketManager.subscribe should subscribe from a topic")
# Given a manager
manager = SocketManager(zmq, context)
# And a couple of sockets
manager.ensure_and_bind('foo', zmq.PUB, 'inproc://test.publisher.2', zmq.POLLOUT)
manager.create('bar', zmq.SUB)
manager.connect('bar', 'inproc://test.publisher.2', zmq.POLLIN)
gevent.spawn(manager.publish_safe, 'foo', 'some-topic', {'some': 'value'})
# Then it should have received
topic, result = next(manager.subscribe('bar'))
topic.should.equal('some-topic')
result.should.equal({'some': 'value'})
