Python BrokerConnection.ensure_connection方法代码示例

# 需要导入模块: from kombu.connection import BrokerConnection [as 别名]
# 或者: from kombu.connection.BrokerConnection import ensure_connection [as 别名]
class Producer(object):
    """ Defines the base Rabbit Producer class
    """
    def __init__(self, config):
        self.connection_info = {
            'hostname': config['host'],
            'port': int(config['port']),
            'userid': config['username'],
            'password': config['password'],
            'virtual_host': config['vhost']
        }
        self.connection = BrokerConnection(
            transport=Transport,
            **self.connection_info
        )
        self.connection.ensure_connection(self.connection_error)
        self.exchange = Exchange(name=config['exchange'], channel=self.connection)
        self.route = config.get('route')
        self._producer = messaging.Producer(self.connection, exchange=self.exchange, routing_key=self.route, auto_declare=False)

    def connection_error(self, exc, interval):
        """Defines action on connection error
        """
        if interval > 10:
            raise IOError("Bad connection to rabbit")
        else:
            logging.error("rabbitmq connection error, retry in %s seconds.  error: %s", interval, exc.strerror if hasattr(exc, 'strerror') else exc)

    def publish(self, payload, route=None):
        """Defines the publish action
        Args:
          payload: a message payload
          route: a message route
        """
        if not route:
            route = self.route
        logging.info("Published %s to %s with a route of %s", json.dumps(payload), self.exchange, route)
        return self._producer.publish(payload, route=route)

# 需要导入模块: from kombu.connection import BrokerConnection [as 别名]
# 或者: from kombu.connection.BrokerConnection import ensure_connection [as 别名]
class Consumer(object):
    """
    Defines the base Rabbit Consumer class
    """
    def __init__(self, config):
        self.connection_info = {
            'hostname': config['host'],
            'port': int(config['port']),
            'userid': config['username'],
            'password': config['password'],
            'virtual_host': config['vhost']
        }
        self.connection = BrokerConnection(
            transport=Transport,
            **self.connection_info
        )
        self.connection.ensure_connection(self.connection_error)
        self.exchange = Exchange(name=config['exchange'], channel=self.connection)
        if isinstance(config['queue_name'], list):
            self.queues = [Queue(name, exchange=self.exchange) for name in config['queue_name']]
        else:
            self.queues = [Queue(config['queue_name'], exchange=self.exchange)]
        self.callbacks = []
        self.add_callback(self.on_message)
        self.prefetch_count = 1

    def consume(self):
        """Consumes all messages on the queues
        """
        self.callbacks.append(self.ack_message)
        c = messaging.Consumer(self.connection, queues=self.queues, callbacks=self.callbacks, auto_declare=False)
        c.qos(prefetch_count=self.prefetch_count)
        with c:
            for _ in eventloop(self.connection, timeout=10, ignore_timeouts=True):
                pass

    def add_callback(self, callback):
        """Registers a new callback to be called upon message receipt
        Args:
          callback: a function that accepts body, message as args
        """
        def _build_body(body):
            """Helper method to build body"""
            if isinstance(body, dict):
                return body

            try:
                json_data = json.loads(body)
                if "path" not in json_data:
                    json_data["path"] = [None, None]
                return json_data
            except Exception as e:
                logging.error("Invalid JSON. Exception: %s", e)
            return {"invalid_json": str(body), 'path': [None, None]}

        self.callbacks.append(lambda body, message: callback(_build_body(body), message))

    def add_queue(self, queue_name):
        """Adds a queue to the consumer
        Args:
          queue_name: name of queue to be bound to consumer
        """
        self.queues.append(Queue(queue_name, exchange=self.exchange))

    def connection_error(self, exc, interval):
        """Defines action on connection error
        """
        if interval > 10:
            raise IOError("Bad connection to rabbit")
        else:
            logging.error("rabbitmq connection error, retry in %s seconds.  error: %s", interval, exc.strerror if hasattr(exc, 'strerror') else exc)

    def on_message(self, body, _):
        """Defines action to take on receipt of message
        Args:
          body: a decoded message body
          message: the message object
        """
        logging.info("Recieved message %s from one of the following:%s", json.dumps(body), self.queues)

    def ack_message(self, body, message):
        """Defines acking the message after processing
        Args:
          body: a decoded message body
          message: the message object
        """
        logging.info("Acking message %s", json.dumps(body))
        message.ack()

# 需要导入模块: from kombu.connection import BrokerConnection [as 别名]
# 或者: from kombu.connection.BrokerConnection import ensure_connection [as 别名]
class test_Connection(unittest.TestCase):

    def setUp(self):
        self.conn = BrokerConnection(port=5672, transport=Transport)

    def test_establish_connection(self):
        conn = self.conn
        conn.connect()
        self.assertTrue(conn.connection.connected)
        self.assertEqual(conn.host, "localhost:5672")
        channel = conn.channel()
        self.assertTrue(channel.open)
        self.assertEqual(conn.drain_events(), "event")
        _connection = conn.connection
        conn.close()
        self.assertFalse(_connection.connected)
        self.assertIsInstance(conn.transport, Transport)

    def test__enter____exit__(self):
        conn = self.conn
        context = conn.__enter__()
        self.assertIs(context, conn)
        conn.connect()
        self.assertTrue(conn.connection.connected)
        conn.__exit__()
        self.assertIsNone(conn.connection)
        conn.close()    # again

    def test_close_survives_connerror(self):

        class _CustomError(Exception):
            pass

        class MyTransport(Transport):
            connection_errors = (_CustomError, )

            def close_connection(self, connection):
                raise _CustomError("foo")

        conn = BrokerConnection(transport=MyTransport)
        conn.connect()
        conn.close()
        self.assertTrue(conn._closed)

    def test_ensure_connection(self):
        self.assertTrue(self.conn.ensure_connection())

    def test_SimpleQueue(self):
        conn = self.conn
        q = conn.SimpleQueue("foo")
        self.assertTrue(q.channel)
        self.assertTrue(q.channel_autoclose)
        chan = conn.channel()
        q2 = conn.SimpleQueue("foo", channel=chan)
        self.assertIs(q2.channel, chan)
        self.assertFalse(q2.channel_autoclose)

    def test_SimpleBuffer(self):
        conn = self.conn
        q = conn.SimpleBuffer("foo")
        self.assertTrue(q.channel)
        self.assertTrue(q.channel_autoclose)
        chan = conn.channel()
        q2 = conn.SimpleBuffer("foo", channel=chan)
        self.assertIs(q2.channel, chan)
        self.assertFalse(q2.channel_autoclose)

    def test__repr__(self):
        self.assertTrue(repr(self.conn))

    def test__reduce__(self):
        x = pickle.loads(pickle.dumps(self.conn))
        self.assertDictEqual(x.info(), self.conn.info())

    def test_channel_errors(self):

        class MyTransport(Transport):
            channel_errors = (KeyError, ValueError)

        conn = BrokerConnection(transport=MyTransport)
        self.assertTupleEqual(conn.channel_errors, (KeyError, ValueError))

    def test_connection_errors(self):

        class MyTransport(Transport):
            connection_errors = (KeyError, ValueError)

        conn = BrokerConnection(transport=MyTransport)
        self.assertTupleEqual(conn.connection_errors, (KeyError, ValueError))

# 需要导入模块: from kombu.connection import BrokerConnection [as 别名]
# 或者: from kombu.connection.BrokerConnection import ensure_connection [as 别名]
class test_Connection(TestCase):

    def setUp(self):
        self.conn = BrokerConnection(port=5672, transport=Transport)

    def test_establish_connection(self):
        conn = self.conn
        conn.connect()
        self.assertTrue(conn.connection.connected)
        self.assertEqual(conn.host, "localhost:5672")
        channel = conn.channel()
        self.assertTrue(channel.open)
        self.assertEqual(conn.drain_events(), "event")
        _connection = conn.connection
        conn.close()
        self.assertFalse(_connection.connected)
        self.assertIsInstance(conn.transport, Transport)

    def test__enter____exit__(self):
        conn = self.conn
        context = conn.__enter__()
        self.assertIs(context, conn)
        conn.connect()
        self.assertTrue(conn.connection.connected)
        conn.__exit__()
        self.assertIsNone(conn.connection)
        conn.close()    # again

    def test_close_survives_connerror(self):

        class _CustomError(Exception):
            pass

        class MyTransport(Transport):
            connection_errors = (_CustomError, )

            def close_connection(self, connection):
                raise _CustomError("foo")

        conn = BrokerConnection(transport=MyTransport)
        conn.connect()
        conn.close()
        self.assertTrue(conn._closed)

    def test_close_when_default_channel(self):
        conn = self.conn
        conn._default_channel = Mock()
        conn._close()
        conn._default_channel.close.assert_called_with()

    def test_close_when_default_channel_close_raises(self):

        class Conn(BrokerConnection):

            @property
            def connection_errors(self):
                return (KeyError, )

        conn = Conn("memory://")
        conn._default_channel = Mock()
        conn._default_channel.close.side_effect = KeyError()

        conn._close()
        conn._default_channel.close.assert_called_with()

    def test_revive_when_default_channel(self):
        conn = self.conn
        defchan = conn._default_channel = Mock()
        conn.revive(Mock())

        defchan.close.assert_called_with()
        self.assertIsNone(conn._default_channel)

    def test_ensure_connection(self):
        self.assertTrue(self.conn.ensure_connection())

    def test_ensure_success(self):
        def publish():
            return "foobar"

        ensured = self.conn.ensure(None, publish)
        self.assertEqual(ensured(), "foobar")

    def test_ensure_failure(self):
        class _CustomError(Exception):
            pass

        def publish():
            raise _CustomError("bar")

        ensured = self.conn.ensure(None, publish)
        with self.assertRaises(_CustomError):
            ensured()

    def test_ensure_connection_failure(self):
        class _ConnectionError(Exception):
            pass

        def publish():
            raise _ConnectionError("failed connection")
#.........这里部分代码省略.........