#!/usr/bin/env python
# -*- coding: utf-8 -*-
#


from amqp.connection import Connection as amqp_connection
from amqp import basic_message


class RabbitMQTest(object):

    def __init__(self, host, port, exchange="test", queue="test", username="test", password="test", virtual_host="/"):

        self.connection = amqp_connection(
                            host=host,
                            port=port,
                            virtual_host=virtual_host,
                            userid=username,
                            password=password
                            )

        self.exchange = exchange
        self.queue = queue
        self.channel = self.connection.channel()

    def createSimpleExchangeQueue(self):

        self.channel.exchange_declare(
            exchange=self.exchange,
            type="direct",
            auto_delete=True

        )

        self.channel.queue_declare(
            queue=self.queue,
            auto_delete=True
        )

        self.channel.queue_bind(
            queue=self.queue,
            exchange=self.exchange
        )

    def dumpMessages(self, number=10240, size=1024):

        print ("Dumping %s messages of %s bytes." % (number, size))
        message = basic_message.Message(
            body="x" * size)

        for _ in range(number):
            self.channel.basic_publish(
                message,
                exchange=self.exchange
            )

    def createFeedbackLoop(self, total_messages):

        print ("Producing and consuming %s messages." % (total_messages))

        self.counter = 0

        def produce(message):

            self.channel.basic_publish(
               message,
               exchange=self.exchange
            )

            self.channel.basic_ack(message.delivery_tag)
            self.counter += 1

        self.channel.basic_qos(
            prefetch_size=0,
            prefetch_count=100,
            a_global=False
        )

        self.channel.basic_consume(
            queue=self.exchange,
            callback=produce,
            no_ack=False)

        while self.counter < total_messages:
            self.connection.drain_events()


def main():

    setup = RabbitMQTest(
        "your-rabbitmq-host-01",
        5672,
        exchange="ha_test",
        queue="ha_test")

    setup.createSimpleExchangeQueue()

    # dump 1000mb worth of messages into the queue
    setup.dumpMessages(number=102400)

    # Consume & produce 1000000 messages
    setup.createFeedbackLoop(1000000)

if __name__ == "__main__":
    main()