def test_socket_manager_subscribe_invalid_callable(get_by_name): ( "SocketManager().subscribe() should raise TypeError " "if the keep_polling callable is not callable" ) # Background: wait_until_ready is mocked to return the socket socket = get_by_name.return_value socket.recv_multipart.return_value = ["metrics:whatevs", "the-data"] # 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 perform one iteration in the subscriber when_called = list.when.called_with( manager.subscribe("foobar", keep_polling="not a callable") ) when_called.should.have.raised( TypeError, "SocketManager.subscribe parameter keep_polling must be a function or callable that returns a boolean", )
def test_socket_manager_set_topic(get_by_name): ("SocketManager().set_topic() should retrieve the socket by " "name and set its subscription topic") # Background: get_by_name is mocked socket = get_by_name.return_value # 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) # When I call set topic manager.set_topic('the-socket-name', 'the-topic-name') # Then it should have retrieved the socket by name get_by_name.assert_called_once_with('the-socket-name') # And the topic should have been set in that socket socket.setsockopt.assert_called_once_with(zmq.SUBSCRIBE, b'the-topic-name')
def test_socket_manager_ready(get_by_name, engage): ("SocketManager().ready() should return the socket " "if it's ready and engaged with the given polling mechanism") # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager with a fake socket and a fake cache manager = SocketManager(zmq, context) manager.sockets = {"foobar": "SOME_TYPE"} manager.cache = {'socket_sentinel': "some-polling-mechanism"} # And that get_by_name will return "socket_sentinel" get_by_name.return_value = 'socket_sentinel' # When I call ready socket = manager.ready('foobar', 'some-polling-mechanism', timeout=1000) # Then the result should be the "socket_sentinel" socket.should.equal("socket_sentinel") # And engage should have been called engage.assert_called_once_with(1000)
def test_socket_bind_to_random_port(get_by_name, register_socket, engage): ("SocketManager().bind_to_random_port() should get by name, register then " "bind to a random port") # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager manager = SocketManager(zmq, context) # And that the socket is mocked to return a random port socket = get_by_name.return_value socket.bind_to_random_port.return_value = 42000 # When I call bind result = manager.bind_to_random_port('foobar', 'some-mechanism') # Then it should have returned a socket from get_by_name result.should.equal((socket, 'tcp://0.0.0.0:42000')) # And get_by_name should have been called with the name get_by_name.assert_called_once_with('foobar') # And register_socket should have been called register_socket.assert_called_once_with(socket, "some-mechanism") # And it should have called bind on the given address socket.bind_to_random_port.assert_called_once_with('tcp://0.0.0.0') # And it should have engaged the manager with zero timeout engage.assert_called_once_with(0)
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_bind(get_by_name, register_socket, engage): ("SocketManager().bind() should get by name, register 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 bind result = manager.bind('foobar', 'inproc://whatevs', 'some-mechanism') # Then it should have returned a socket from get_by_name result.should.equal(get_by_name.return_value) # And get_by_name should have been called with the name get_by_name.assert_called_once_with('foobar') # And register_socket should have been called register_socket.assert_called_once_with(result, "some-mechanism") # And it should have called bind on the given address result.bind.assert_called_once_with('inproc://whatevs') # And it should have engaged the manager with zero timeout engage.assert_called_once_with(0)
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_raises_exception_when_creating_existing_socket(): ("SocketManager().create() should raise AgentZeroSocketError when " "a given socket already exists") # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager manager = SocketManager(zmq, context) # And a couple of sockets manager.create('foo', zmq.REQ) manager.create('bar', zmq.PUB) # When I call create on an existing socket name when_called = manager.create.when.called_with('foo', zmq.PUB) # Then it should have when_called.should.have.raised( SocketAlreadyExists, "SocketManager(sockets=['foo', 'bar']) already has a socket named 'foo'" )
def test_socket_manager_raises_exception_when_retrieving_socket_with_bad_name( ): ("SocketManager().get_by_name() should raise AgentZeroSocketError when " "a given socket does not exist") # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager manager = SocketManager(zmq, context) # And a couple of sockets manager.create('foo', zmq.REQ) manager.create('bar', zmq.PUB) # When I call get_by_name with an unexpected name when_called = manager.get_by_name.when.called_with('boom') # Then it should have when_called.should.have.raised( SocketNotFound, "SocketManager(sockets=['foo', 'bar']) has no sockets named 'boom'.")
def test_socket_manager_get_or_create_preexisting(create, get_by_name, register_socket): ("SocketManager().get_or_create() should just return " "and register a pre-existing socket with the same name") # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager with no sockets manager = SocketManager(zmq, context) manager.sockets = {} # When I call get_or_create result = manager.get_or_create("foobar", "SOME_TYPE", "fake-polling-mechanism") # Then it should have returned a socket result.should.equal(get_by_name.return_value) # And the socket should have been created create.assert_called_once_with("foobar", "SOME_TYPE") # And should have been retrieved by name get_by_name.assert_called_once_with("foobar") # And should have been registered register_socket.assert_called_once_with(result, "fake-polling-mechanism")
def test_socket_manager_not_ready(get_by_name, engage): ("SocketManager().ready() should return None when the socket is not available" ) # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager with a fake socket and an *EMPTY* cache manager = SocketManager(zmq, context) manager.sockets = {"foobar": "SOME_TYPE"} manager.cache = {} # And that get_by_name will return "socket_sentinel" get_by_name.return_value = 'socket_sentinel' # When I call ready socket = manager.ready('foobar', 'some-polling-mechanism', timeout=1000) # Then the result should be None socket.should.be.none # And engage should have been called engage.assert_called_once_with(1000)
def test_socket_connect(get_by_name, register_socket, engage): ( "SocketManager().connect() should get by name, register then " "connect 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 connect result = manager.connect("foobar", "inproc://whatevs", "some-mechanism") # Then it should have returned a socket from get_by_name result.should.equal(get_by_name.return_value) # And get_by_name should have been called with the name get_by_name.assert_called_once_with("foobar") # And register_socket should have been called register_socket.assert_called_once_with(result, "some-mechanism") # And it should have called connect on the given address result.connect.assert_called_once_with("inproc://whatevs") # And it should have engaged the manager with zero timeout engage.assert_called_once_with(0)
def test_socket_manager_recv_event_safe(wait_until_ready, set_topic): ("SocketManager().recv_event_safe() should set the topic " "on the given socket, wait for data to become available " "and return it") # Background: wait_until_ready is mocked to return the socket wait_until_ready.side_effect = lambda name, *args: manager.sockets[name] # 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) socket.recv_multipart.return_value = b'a fake topic', 'the-raw-payload' # When I call .recv_event_safe() result = manager.recv_event_safe('foobar', topic=b'helloworld') # Then it should have unpacked the payload after receiving serializer.unpack.assert_called_once_with('the-raw-payload') # And the result should be the unpacked value result.should.be.an(Event)
def test_socket_manager_send_safe_not_ready(wait_until_ready): ("SocketManager().send_safe() should return False when the socet is not ready") # Background: wait_until_ready is mocked to return None wait_until_ready.return_value = None # 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) # When I call .send_safe() sent = manager.send_safe('foobar', 'PAYLOAD') # Then it should have failed sent.should.be.false # And it should not pack the value serializer.pack.called.should.be.false
def test_socket_manager_ready(get_by_name, engage): ("SocketManager().ready() should return the socket " "if it's ready and engaged with the given polling mechanism") # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager with a fake socket and a fake cache manager = SocketManager(zmq, context) manager.sockets = { "foobar": "SOME_TYPE" } manager.cache = { 'socket_sentinel': "some-polling-mechanism" } # And that get_by_name will return "socket_sentinel" get_by_name.return_value = 'socket_sentinel' # When I call ready socket = manager.ready('foobar', 'some-polling-mechanism', timeout=1000) # Then the result should be the "socket_sentinel" socket.should.equal("socket_sentinel") # And engage should have been called engage.assert_called_once_with(1000)
def test_socket_engage(): ("SocketManager().engage() should return an OrderedDict with the polled sockets") # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager manager = SocketManager(zmq, context) # And that the poller returns some results manager.poller.poll.return_value = [("foobar", "whatevs"), ("silly-walks", "whatevs")] # When I call engage result = manager.engage("SOME_TIMEOUT") # Then it should have called .poll() with the given timeout manager.poller.poll.assert_called_once_with("SOME_TIMEOUT") # And the result should be an OrderedDict result.should.be.an(OrderedDict) # And the result should contain the returned items from polling result.items().should.equal([("foobar", "whatevs"), ("silly-walks", "whatevs")])
def test_socket_manager_get_or_create_preexisting(create, get_by_name, register_socket): ("SocketManager().get_or_create() should just return " "and register a pre-existing socket with the same name") # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager with no sockets manager = SocketManager(zmq, context) manager.sockets = {} # When I call get_or_create result = manager.get_or_create("foobar", "SOME_TYPE", "fake-polling-mechanism") # Then it should have returned a socket result.should.equal(get_by_name.return_value) # And the socket should have been created create.assert_called_once_with("foobar", "SOME_TYPE") # And should have been retrieved by name get_by_name.assert_called_once_with("foobar") # And should have been registered register_socket.assert_called_once_with( result, "fake-polling-mechanism" )
def test_socket_manager_raises_exception_when_creating_existing_socket(): ("SocketManager().create() should raise AgentZeroSocketError when " "a given socket already exists") # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager manager = SocketManager(zmq, context) # And a couple of sockets manager.create('foo', zmq.REQ) manager.create('bar', zmq.PUB) # When I call create on an existing socket name when_called = manager.create.when.called_with('foo', zmq.PUB) # Then it should have when_called.should.have.raised( SocketAlreadyExists, "SocketManager(sockets=['foo', 'bar']) already has a socket named 'foo'" )
def test_socket_engage(): ("SocketManager().engage() should return an OrderedDict with the polled sockets" ) # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager manager = SocketManager(zmq, context) # And that the poller returns some results manager.poller.poll.return_value = [("foobar", "whatevs"), ("silly-walks", "whatevs")] # When I call engage result = manager.engage("SOME_TIMEOUT") # Then it should have called .poll() with the given timeout manager.poller.poll.assert_called_once_with("SOME_TIMEOUT") # And the result should be an OrderedDict result.should.be.an(OrderedDict) # And the result should contain the returned items from polling result.items().should.equal([("foobar", "whatevs"), ("silly-walks", "whatevs")])
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_manager_subscribe(get_by_name): ("SocketManager().subscribe() should yield the topic and ") # Background: wait_until_ready is mocked to return the socket socket = get_by_name.return_value socket.recv_multipart.return_value = ['metrics:whatevs', 'the-data'] # 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 perform one iteration in the subscriber topic, payload = next(manager.subscribe('foobar')) # Then it should have unpacked the payload after receiving serializer.unpack.assert_called_once_with('the-data') # And the result should be the unpacked value payload.should.equal(serializer.unpack.return_value) # And the topic should be the expected topic.should.equal('metrics:whatevs')
def test_socket_manager_recv_event_safe(wait_until_ready, set_topic): ("SocketManager().recv_event_safe() should set the topic " "on the given socket, wait for data to become available " "and return it") # Background: wait_until_ready is mocked to return the socket wait_until_ready.side_effect = lambda name, *args: manager.sockets[name] # 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) socket.recv_multipart.return_value = b'a fake topic', 'the-raw-payload' # When I call .recv_event_safe() result = manager.recv_event_safe('foobar', topic=b'helloworld') # Then it should have unpacked the payload after receiving serializer.unpack.assert_called_once_with('the-raw-payload') # And the result should be the unpacked value result.should.be.an(Event)
def test_socket_manager_recv_safe_not_ready(wait_until_ready): ("SocketManager().recv_safe() should return None when the socket is not ready" ) # Background: wait_until_ready is mocked to return None wait_until_ready.return_value = None # 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) # When I call .recv_safe() result = manager.recv_safe('foobar') # Then it should never had anything to unpack serializer.unpack.called.should.be.false # And the result should be None result.should.be.none
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_recv_safe(wait_until_ready): ("SocketManager().recv_safe() should deserialize " "after receiving, using the configured backend") # Background: wait_until_ready is mocked to return the socket wait_until_ready.side_effect = lambda name, *args: manager.sockets[name] # 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 .recv_safe() result = manager.recv_safe('foobar') # Then it should have unpacked the payload after receiving serializer.unpack.assert_called_once_with(socket.recv.return_value) # And the result should be the unpacked value result.should.equal(serializer.unpack.return_value)
def test_socket_manager_send_safe_not_ready(wait_until_ready): ("SocketManager().send_safe() should return False when the socet is not ready" ) # Background: wait_until_ready is mocked to return None wait_until_ready.return_value = None # 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) # When I call .send_safe() sent = manager.send_safe('foobar', 'PAYLOAD') # Then it should have failed sent.should.be.false # And it should not pack the value serializer.pack.called.should.be.false
def test_socket_manager_send_safe(wait_until_ready): ("SocketManager().send_safe() should serialize " "before sending, using the configured backend") # Background: wait_until_ready is mocked to return the socket wait_until_ready.side_effect = lambda name, *args: manager.sockets[name] # 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 .send_safe() sent = manager.send_safe('foobar', 'PAYLOAD') # Then it should have sent successfully sent.should.be.true # And it should have packed the payload before sending serializer.pack.assert_called_once_with('PAYLOAD') packed = serializer.pack.return_value socket.send.assert_called_once_with(packed)
def test_socket_manager_recv_safe_not_ready(wait_until_ready): ("SocketManager().recv_safe() should return None when the socket is not ready") # Background: wait_until_ready is mocked to return None wait_until_ready.return_value = None # 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) # When I call .recv_safe() result = manager.recv_safe('foobar') # Then it should never had anything to unpack serializer.unpack.called.should.be.false # And the result should be None result.should.be.none
def test_socket_manager_recv_event_safe_missing_socket(wait_until_ready, set_topic): ("SocketManager().recv_event_safe() should return None when the socket is not ready" ) # Background: wait_until_ready is mocked to return None wait_until_ready.return_value = None # 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) socket.recv_multipart.return_value = b'a fake topic', 'the-raw-payload' # When I call .recv_event_safe() result = manager.recv_event_safe('foobar', topic=b'helloworld') # And the result should be None result.should.be.none
def test_socket_manager_subscribe(get_by_name): ("SocketManager().subscribe() should yield the topic and ") # Background: wait_until_ready is mocked to return the socket socket = get_by_name.return_value socket.recv_multipart.return_value = ['metrics:whatevs', 'the-data'] # 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 perform one iteration in the subscriber topic, payload = next(manager.subscribe('foobar')) # Then it should have unpacked the payload after receiving serializer.unpack.assert_called_once_with('the-data') # And the result should be the unpacked value payload.should.equal(serializer.unpack.return_value) # And the topic should be the expected topic.should.equal('metrics:whatevs')
def test_socket_manager_recv_event_safe_no_topic(wait_until_ready, set_topic): ("SocketManager().recv_event_safe() should raise an exeption when the topic is not a string" ) # Background: wait_until_ready is mocked to return the socket wait_until_ready.side_effect = lambda name, *args: manager.sockets[name] # 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) socket.recv_multipart.return_value = b'a fake topic', 'the-raw-payload' # When I call .recv_event_safe() when_called = manager.recv_event_safe.when.called_with('foobar', topic={'boom'}) # Then it should have raised and exception when_called.should.have.raised( TypeError, "recv_event_safe() takes a string, None " "or False as argument, received " "set(['boom'])(<type 'set'>) instead")
def test_socket_manager_register_socket(): ("SocketManager().register_socket() should register " "a socket only once even if called twice") # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager with a fake socket in the pool manager = SocketManager(zmq, context) # When I call register_socket result = manager.register_socket("fake-socket", "fake-polling-mechanism") # Then it should return the socket itself result.should.equal("fake-socket") # And it should have been registered with the poller manager.poller.register.assert_called_once_with( "fake-socket", "fake-polling-mechanism", ) # And it should have been added to the internal registry manager.registry.should.have.key("fake-socket").being.equal("fake-polling-mechanism")
def test_socket_disconnect(get_by_name, register_socket, engage): ( "SocketManager().disconnect() should get by name, unregister then " "disconnect to the given address." ) socket = get_by_name.return_value # Given a zmq mock zmq = Mock() poller = zmq.Poller.return_value # And a poller that raises an exception upon unregister poller.unregister.side_effect = RuntimeError("boom") # And a context context = Mock() # And a socket manager manager = SocketManager(zmq, context) manager.addresses = {"foobar": "baz", "another": "socket"} # When I call disconnect manager.disconnect("foobar").should.be.true # Then it should have removed the socket address from the table manager.addresses.should.equal({"another": "socket"}) socket.disconnect.assert_has_calls([call("baz")]) socket.disconnect.call_count.should.equal(1)
def test_socket_manager_get_or_create(create, get_by_name, register_socket): ("SocketManager().get_or_create() should create a socket and register it") # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager with no sockets manager = SocketManager(zmq, context) manager.sockets = { "foobar": "SOME_TYPE" } # When I call get_or_create result = manager.get_or_create("foobar", "SOME_TYPE", "fake-polling-mechanism") # Then it should have returned a socket result.should.equal(get_by_name.return_value) # And .create should not have been called create.called.should.be.false # And should have been retrieved by name get_by_name.assert_called_once_with("foobar") # And should have been registered register_socket.assert_called_once_with( result, "fake-polling-mechanism", )
def test_socket_manager_recv_event_safe_missing_socket(wait_until_ready, set_topic): ("SocketManager().recv_event_safe() should return None when the socket is not ready") # Background: wait_until_ready is mocked to return None wait_until_ready.return_value = None # 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) socket.recv_multipart.return_value = b'a fake topic', 'the-raw-payload' # When I call .recv_event_safe() result = manager.recv_event_safe('foobar', topic=b'helloworld') # And the result should be None result.should.be.none
def test_socket_manager_raises_exception_when_retrieving_socket_with_bad_name(): ("SocketManager().get_by_name() should raise AgentZeroSocketError when " "a given socket does not exist") # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager manager = SocketManager(zmq, context) # And a couple of sockets manager.create('foo', zmq.REQ) manager.create('bar', zmq.PUB) # When I call get_by_name with an unexpected name when_called = manager.get_by_name.when.called_with('boom') # Then it should have when_called.should.have.raised( SocketNotFound, "SocketManager(sockets=['foo', 'bar']) has no sockets named 'boom'." )
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_bind(get_by_name, register_socket, engage): ("SocketManager().bind() should get by name, register 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 bind result = manager.bind('foobar', 'inproc://whatevs', 'some-mechanism') # Then it should have returned a socket from get_by_name result.should.equal(get_by_name.return_value) # And get_by_name should have been called with the name get_by_name.assert_called_once_with('foobar') # And register_socket should have been called register_socket.assert_called_once_with( result, "some-mechanism" ) # And it should have called bind on the given address result.bind.assert_called_once_with('inproc://whatevs') # And it should have engaged the manager with zero timeout engage.assert_called_once_with(0)
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_not_ready(get_by_name, engage): ("SocketManager().ready() should return None when the socket is not available") # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager with a fake socket and an *EMPTY* cache manager = SocketManager(zmq, context) manager.sockets = { "foobar": "SOME_TYPE" } manager.cache = {} # And that get_by_name will return "socket_sentinel" get_by_name.return_value = 'socket_sentinel' # When I call ready socket = manager.ready('foobar', 'some-polling-mechanism', timeout=1000) # Then the result should be None socket.should.be.none # And engage should have been called engage.assert_called_once_with(1000)
def test_socket_manager_register_socket(): ("SocketManager().register_socket() should register " "a socket only once even if called twice") # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager with a fake socket in the pool manager = SocketManager(zmq, context) # When I call register_socket result = manager.register_socket("fake-socket", "fake-polling-mechanism") # Then it should return the socket itself result.should.equal("fake-socket") # And it should have been registered with the poller manager.poller.register.assert_called_once_with( "fake-socket", "fake-polling-mechanism", ) # And it should have been added to the internal registry manager.registry.should.have.key("fake-socket").being.equal( "fake-polling-mechanism")
def test_socket_manager_send_safe(wait_until_ready): ("SocketManager().send_safe() should serialize " "before sending, using the configured backend") # Background: wait_until_ready is mocked to return the socket wait_until_ready.side_effect = lambda name, *args: manager.sockets[name] # 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 .send_safe() sent = manager.send_safe('foobar', 'PAYLOAD') # Then it should have sent successfully sent.should.be.true # And it should have packed the payload before sending serializer.pack.assert_called_once_with('PAYLOAD') packed = serializer.pack.return_value socket.send.assert_called_once_with(packed)
def test_socket_manager_get_or_create(create, get_by_name, register_socket): ("SocketManager().get_or_create() should create a socket and register it") # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager with no sockets manager = SocketManager(zmq, context) manager.sockets = {"foobar": "SOME_TYPE"} # When I call get_or_create result = manager.get_or_create("foobar", "SOME_TYPE", "fake-polling-mechanism") # Then it should have returned a socket result.should.equal(get_by_name.return_value) # And .create should not have been called create.called.should.be.false # And should have been retrieved by name get_by_name.assert_called_once_with("foobar") # And should have been registered register_socket.assert_called_once_with( result, "fake-polling-mechanism", )
def test_socket_manager_recv_safe(wait_until_ready): ("SocketManager().recv_safe() should deserialize " "after receiving, using the configured backend") # Background: wait_until_ready is mocked to return the socket wait_until_ready.side_effect = lambda name, *args: manager.sockets[name] # 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 .recv_safe() result = manager.recv_safe('foobar') # Then it should have unpacked the payload after receiving serializer.unpack.assert_called_once_with(socket.recv.return_value) # And the result should be the unpacked value result.should.equal(serializer.unpack.return_value)
def test_socket_manager_recv_event_safe_no_topic(wait_until_ready, set_topic): ("SocketManager().recv_event_safe() should raise an exeption when the topic is not a string") # Background: wait_until_ready is mocked to return the socket wait_until_ready.side_effect = lambda name, *args: manager.sockets[name] # 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) socket.recv_multipart.return_value = b'a fake topic', 'the-raw-payload' # When I call .recv_event_safe() when_called = manager.recv_event_safe.when.called_with('foobar', topic={'boom'}) # Then it should have raised and exception when_called.should.have.raised( TypeError, "recv_event_safe() takes a string, None " "or False as argument, received " "set(['boom'])(<type 'set'>) instead" )
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_set_topic(get_by_name): ("SocketManager().set_topic() should retrieve the socket by " "name and set its subscription topic") # Background: get_by_name is mocked socket = get_by_name.return_value # 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) # When I call set topic manager.set_topic('the-socket-name', 'the-topic-name') # Then it should have retrieved the socket by name get_by_name.assert_called_once_with('the-socket-name') # And the topic should have been set in that socket socket.setsockopt.assert_called_once_with( zmq.SUBSCRIBE, b'the-topic-name' )
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()
class GraphClient(object): def __init__(self, request_connect_addr='tcp://127.0.0.1:6000'): self.context = zmq.Context() self.sockets = SocketManager(zmq, self.context) self.sockets.create('query', zmq.REQ) def request(self, data): self.sockets.send_safe('query', data) return self.sockets.recv_safe('query')
class Pipeline(object): steps = [] def __init__(self): self.context = zmq.Context() self.sockets = SocketManager(zmq, self.context) self.sockets.create("pipe-sub", zmq.SUB) self.sockets.create("pipe-in", zmq.PULL) self.sockets.create("pipe-out", zmq.PUSH) self.children = []
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'})
def test_socket_manager_send_safe_not_ready(): ("SocketManager.send_safe should return False when the socket is not ready") # Given a manager manager = SocketManager(zmq, context) # And a couple of sockets manager.create('foo', zmq.REP) # When I call .send_safe() result = manager.send_safe('foo', {'some': 'value'}) # Then it should be false result.should.be.false
def test_socket_manager_ignore_exceptions_when_cleaning_up_sockets(): ("SocketManager.__del__() should should ignore errors when cleaning up sockets") # Given a bad socket that raises an exception upon calling .close() class BadSocket(object): def close(self): raise RuntimeError("This is the exception that should be ignored") # And a SocketManager instance containing that poisened socket manager = SocketManager(zmq, context) manager.sockets["bad-one"] = BadSocket() # When the manager gets wipe out of the memory del manager
class GraphClient(object): def __init__(self, request_connect_addr='tcp://127.0.0.1:6000'): self.context = zmq.Context() self.sockets = SocketManager(zmq, self.context) self.sockets.ensure_and_connect('query', zmq.REQ, request_connect_addr, zmq.POLLIN | zmq.POLLOUT) print 'connected to', request_connect_addr def request(self, data): self.sockets.send_safe('query', data) result = self.sockets.recv_safe('query') return result def query(self, string): return self.request({'query': string})
def test_socket_manager_init(): ("SocketManager() should set an internal zmq module and create a poller") # Given a zmq mock zmq = Mock() # And a context context = Mock() # When I create a socket manager manager = SocketManager(zmq, context) # Then it should have set the zmq module and created a poller manager.should.have.property('zmq').being.equal(zmq) manager.should.have.property('poller').being.equal(zmq.Poller.return_value) # And it should have an OrderedDict to keep track of the socket names manager.should.have.property('sockets').being.an(OrderedDict) manager.sockets.should.be.empty # And it should have an OrderedDict to keep track of the registered sockets manager.should.have.property('registry').being.an(OrderedDict) manager.registry.should.be.empty # And it should have a dict for caching polled sockets manager.should.have.property('cache').being.a(dict) manager.cache.should.be.empty
def test_socket_manager_wait_until_ready(ready, engage, gevent, time): ("SocketManager().wait_until_ready() whould engage and poll " "continously until the socket becomes available") # Background: time.time() is mocked tick for 3 iterations time.time.side_effect = range(3) # Background: SocketManager().ready() is mocked to return a socket # only after the second iteration socket = Mock(name='socket-ready-to-go') ready.side_effect = [None, socket] # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager manager = SocketManager(zmq, context) # When I call wait_until_ready result = manager.wait_until_ready("foobar", "some-mechanism", timeout=2, polling_timeout=42) # Then it should have returned the socket result.should.equal(socket) # And it should have engaged twice engage.assert_has_calls([ call(42), call(42), ]) # And ready should also have been called twice ready.assert_has_calls([ call('foobar', 'some-mechanism'), call('foobar', 'some-mechanism'), ]) # And gevent.sleep should have been called twice gevent.sleep.assert_has_calls([ call(), call(), ])
def test_socket_manager_repr(): ("repr(SocketManager()) should be informative") # Given a zmq mock zmq = Mock() # And a context context = Mock() # And a socket manager with some fake sockets manager = SocketManager(zmq, context) manager.sockets = OrderedDict([("foobar", "WHATEVS"), ("awesome", "WHATEVS")]) # When I get the string representation of the instance string = repr(manager) # Then it should string.should.equal("SocketManager(sockets=['foobar', 'awesome'])")