Python core.SocketManager类代码示例

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