Python SocketManager.recv_event_safe方法代码示例

# 需要导入模块: from agentzero.core import SocketManager [as 别名]
# 或者: from agentzero.core.SocketManager import recv_event_safe [as 别名]
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

# 需要导入模块: from agentzero.core import SocketManager [as 别名]
# 或者: from agentzero.core.SocketManager import recv_event_safe [as 别名]
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)

# 需要导入模块: from agentzero.core import SocketManager [as 别名]
# 或者: from agentzero.core.SocketManager import recv_event_safe [as 别名]

#.........这里部分代码省略.........
    def bind_action(self, name, method=None):
        action = getattr(self.actions, name, None)
        if not action:
            raise KeyError('undefined action: {0}'.format(name))

        method = method or getattr(self, 'on_{0}'.format(name), None)
        if not method:
            raise TypeError('{0} does not have method {1}(self, topic, data)'.format(self.__class__, name))

        action(lambda _, event: self.spawn(method, event))

    def should_run(self):
        return self._allowed_to_run

    def listen(self, subscriber_bind_address='tcp://127.0.0.1:6000', pull_bind_address='tcp://127.0.0.1:7000'):
        self.sockets.bind('step-events', subscriber_bind_address, zmq.POLLIN)
        self.sockets.bind('jobs-in', pull_bind_address, zmq.POLLIN)
        self.logger.info('listening for events on %s', subscriber_bind_address)
        self.logger.info('listening for instructions on %s', pull_bind_address)

    def route_event(self, event):
        if not event:
            return

        ROUTES = {
            re.compile(r'available'): self.actions.available,
            re.compile(r'failed'): self.actions.failed,
            re.compile(r'success'): self.actions.success,
            re.compile(r'started'): self.actions.started,
            re.compile(r'metric'): self.actions.metric,
            re.compile(r'logs'): self.actions.logs,
            re.compile(r'error'): self.actions.error,
        }
        matched = False
        for regex, action in ROUTES.items():
            if regex.search(event.topic):
                action.shout(event)
                matched = True

        if not matched:
            print 'unmatched event', event.topic, event.data

    def drain_jobs_in(self):
        while self.should_run():
            data = self.sockets.recv_safe('jobs-in')
            if not data:
                gevent.sleep(0)
                continue

            job = Job.new(data)
            self.backend.enqueue_job(job)
            gevent.sleep(0)

    def drain_jobs_out(self):
        iteration = -1
        while self.should_run():
            iteration += 1

            index = iteration % len(self.steps)

            job_type = self.steps[index]
            worker = self.backend.get_next_available_worker_for_type(job_type)

            if not worker:
                gevent.sleep(0)
                continue

            job = self.backend.dequeue_job_of_type(job_type)
            if not job:
                gevent.sleep(0)
                continue

            self.sockets.send_safe(worker.job_type, job.to_dict())
            gevent.sleep(0)

    def spawn(self, *args, **kw):
        self.greenlets.append(
            self.pool.spawn(*args, **kw)
        )

    def idle(self):
        gevent.sleep(0)

    def loop(self):
        self.listen()
        self.spawn(self.drain_events)
        self.spawn(self.drain_jobs_in)
        self.spawn(self.drain_jobs_out)
        while self.should_run():
            gevent.sleep(5)

    def drain_events(self):
        # drain events
        while self.should_run():
            event = self.sockets.recv_event_safe('step-events')
            if event:
                self.route_event(event)
                gevent.sleep(0)
            else:
                self.idle()