Merge pull request #23 from ramielrowe/queue_args

Worker tests and queue args (for HA support)

.gitignore

@@ -2,3 +2,5 @@
 
 *.pyc
 local_settings.py
+.coverage
+worker.log

run_tests.sh

@@ -1,2 +1,2 @@
 #!/bin/bash
-nosetests tests --exclude-dir=stacktach --with-coverage --cover-package=stacktach
+nosetests tests --exclude-dir=stacktach --with-coverage --cover-package=stacktach,worker

stacktach/db.py

@@ -1,5 +1,8 @@
 import models
 
+def get_or_create_deployment(name):
+    return models.Deployment.objects.get_or_create(name=name)
+
 def create_rawdata(**kwargs):
     return models.RawData(**kwargs)
 

stacktach/models.py

@@ -21,10 +21,6 @@ class Deployment(models.Model):
     name = models.CharField(max_length=50)
 
 
-def get_or_create_deployment(name):
-    return Deployment.objects.get_or_create(name=name)
-
-
 class RawData(models.Model):
     deployment = models.ForeignKey(Deployment)
     tenant = models.CharField(max_length=50, null=True, blank=True,

tests/unit/test_worker.py

@@ -0,0 +1,189 @@
+import json
+import unittest
+
+import kombu
+import kombu.entity
+import kombu.connection
+import mox
+
+from stacktach import db, views
+import worker.worker as worker
+
+class NovaConsumerTestCase(unittest.TestCase):
+    def setUp(self):
+        self.mox = mox.Mox()
+
+    def tearDown(self):
+        self.mox.UnsetStubs()
+
+    def test_get_consumers(self):
+        created_queues = []
+        created_callbacks =  []
+        created_consumers = []
+        def Consumer(queues=None, callbacks=None):
+            created_queues.extend(queues)
+            created_callbacks.extend(callbacks)
+            consumer = self.mox.CreateMockAnything()
+            created_consumers.append(consumer)
+            return consumer
+        self.mox.StubOutWithMock(worker.NovaConsumer, '_create_exchange')
+        self.mox.StubOutWithMock(worker.NovaConsumer, '_create_queue')
+        consumer = worker.NovaConsumer('test', None, None, True, {})
+        exchange = self.mox.CreateMockAnything()
+        consumer._create_exchange('nova', 'topic').AndReturn(exchange)
+        info_queue = self.mox.CreateMockAnything()
+        error_queue = self.mox.CreateMockAnything()
+        consumer._create_queue('monitor.info', exchange, 'monitor.info')\
+                .AndReturn(info_queue)
+        consumer._create_queue('monitor.error', exchange, 'monitor.error')\
+                .AndReturn(error_queue)
+        self.mox.ReplayAll()
+        consumers = consumer.get_consumers(Consumer, None)
+        self.assertEqual(len(consumers), 1)
+        self.assertEqual(consumers[0], created_consumers[0])
+        self.assertEqual(len(created_queues), 2)
+        self.assertTrue(info_queue in created_queues)
+        self.assertTrue(error_queue in created_queues)
+        self.assertEqual(len(created_callbacks), 1)
+        self.assertTrue(consumer.on_nova in created_callbacks)
+        self.mox.VerifyAll()
+
+    def test_create_exchange(self):
+        args = {'key': 'value'}
+        consumer = worker.NovaConsumer('test', None, None, True, args)
+
+        self.mox.StubOutClassWithMocks(kombu.entity, 'Exchange')
+        exchange = kombu.entity.Exchange('nova', type='topic', exclusive=False,
+                                         durable=True, auto_delete=False)
+        self.mox.ReplayAll()
+        actual_exchange = consumer._create_exchange('nova', 'topic')
+        self.assertEqual(actual_exchange, exchange)
+        self.mox.VerifyAll()
+
+    def test_create_queue(self):
+        self.mox.StubOutClassWithMocks(kombu, 'Queue')
+        exchange = self.mox.CreateMockAnything()
+        queue = kombu.Queue('name', exchange, auto_delete=False, durable=True,
+                            exclusive=False, routing_key='routing.key',
+                            queue_arguments={})
+        consumer = worker.NovaConsumer('test', None, None, True, {})
+        self.mox.ReplayAll()
+        actual_queue = consumer._create_queue('name', exchange, 'routing.key',
+                                              exclusive=False,
+                                              auto_delete=False)
+        self.assertEqual(actual_queue, queue)
+        self.mox.VerifyAll()
+
+
+
+    def test_create_queue_with_queue_args(self):
+        self.mox.StubOutClassWithMocks(kombu, 'Queue')
+        exchange = self.mox.CreateMockAnything()
+        queue_args = {'key': 'value'}
+        queue = kombu.Queue('name', exchange, auto_delete=False, durable=True,
+                            exclusive=False, routing_key='routing.key',
+                            queue_arguments=queue_args)
+        consumer = worker.NovaConsumer('test', None, None, True, queue_args)
+        self.mox.ReplayAll()
+        actual_queue = consumer._create_queue('name', exchange, 'routing.key',
+                                              exclusive=False,
+                                              auto_delete=False)
+        self.assertEqual(actual_queue, queue)
+        self.mox.VerifyAll()
+
+    def test_process(self):
+        deployment = self.mox.CreateMockAnything()
+        raw = self.mox.CreateMockAnything()
+        message = self.mox.CreateMockAnything()
+
+        consumer = worker.NovaConsumer('test', None, deployment, True, {})
+        routing_key = 'monitor.info'
+        message.delivery_info = {'routing_key': routing_key}
+        body_dict = {u'key': u'value'}
+        message.body = json.dumps(body_dict)
+        self.mox.StubOutWithMock(views, 'process_raw_data',
+                                 use_mock_anything=True)
+        args = (routing_key, body_dict)
+        views.process_raw_data(deployment, args, json.dumps(args))\
+             .AndReturn(raw)
+        self.mox.StubOutWithMock(consumer, '_check_memory',
+                                use_mock_anything=True)
+        consumer._check_memory()
+        self.mox.ReplayAll()
+        consumer._process(message)
+        self.assertEqual(consumer.processed, 1)
+        self.mox.VerifyAll()
+
+    def test_run(self):
+        config = {
+            'name': 'east_coast.prod.global',
+            'durable_queue': False,
+            'rabbit_host': '10.0.0.1',
+            'rabbit_port': 5672,
+            'rabbit_userid': 'rabbit',
+            'rabbit_password': 'rabbit',
+            'rabbit_virtual_host': '/'
+        }
+        self.mox.StubOutWithMock(db, 'get_or_create_deployment')
+        deployment = self.mox.CreateMockAnything()
+        db.get_or_create_deployment(config['name'])\
+          .AndReturn((deployment, True))
+        self.mox.StubOutWithMock(kombu.connection, 'BrokerConnection')
+        params = dict(hostname=config['rabbit_host'],
+                      port=config['rabbit_port'],
+                      userid=config['rabbit_userid'],
+                      password=config['rabbit_password'],
+                      transport="librabbitmq",
+                      virtual_host=config['rabbit_virtual_host'])
+        self.mox.StubOutWithMock(worker, "continue_running")
+        worker.continue_running().AndReturn(True)
+        conn = self.mox.CreateMockAnything()
+        kombu.connection.BrokerConnection(**params).AndReturn(conn)
+        conn.__enter__().AndReturn(conn)
+        conn.__exit__(None, None, None).AndReturn(None)
+        self.mox.StubOutClassWithMocks(worker, 'NovaConsumer')
+        consumer = worker.NovaConsumer(config['name'], conn, deployment,
+                                       config['durable_queue'], {})
+        consumer.run()
+        worker.continue_running().AndReturn(False)
+        self.mox.ReplayAll()
+        worker.run(config)
+        self.mox.VerifyAll()
+
+    def test_run_queue_args(self):
+        config = {
+            'name': 'east_coast.prod.global',
+            'durable_queue': False,
+            'rabbit_host': '10.0.0.1',
+            'rabbit_port': 5672,
+            'rabbit_userid': 'rabbit',
+            'rabbit_password': 'rabbit',
+            'rabbit_virtual_host': '/',
+            'queue_arguments': {'x-ha-policy': 'all'}
+        }
+        self.mox.StubOutWithMock(db, 'get_or_create_deployment')
+        deployment = self.mox.CreateMockAnything()
+        db.get_or_create_deployment(config['name'])\
+          .AndReturn((deployment, True))
+        self.mox.StubOutWithMock(kombu.connection, 'BrokerConnection')
+        params = dict(hostname=config['rabbit_host'],
+                      port=config['rabbit_port'],
+                      userid=config['rabbit_userid'],
+                      password=config['rabbit_password'],
+                      transport="librabbitmq",
+                      virtual_host=config['rabbit_virtual_host'])
+        self.mox.StubOutWithMock(worker, "continue_running")
+        worker.continue_running().AndReturn(True)
+        conn = self.mox.CreateMockAnything()
+        kombu.connection.BrokerConnection(**params).AndReturn(conn)
+        conn.__enter__().AndReturn(conn)
+        conn.__exit__(None, None, None).AndReturn(None)
+        self.mox.StubOutClassWithMocks(worker, 'NovaConsumer')
+        consumer = worker.NovaConsumer(config['name'], conn, deployment,
+                                       config['durable_queue'],
+                                       config['queue_arguments'])
+        consumer.run()
+        worker.continue_running().AndReturn(False)
+        self.mox.ReplayAll()
+        worker.run(config)
+        self.mox.VerifyAll()

worker/worker.py

@@ -19,7 +19,6 @@
 import datetime
 import json
 import kombu
-import kombu.connection
 import kombu.entity
 import kombu.mixins
 import logging
@@ -27,8 +26,7 @@ import time
 
 from pympler.process import ProcessMemoryInfo
 
-from stacktach import models, views
+from stacktach import db, views
-from stacktach import datetime_to_decimal as dt
 
 
 LOG = logging.getLogger(__name__)
@@ -39,36 +37,40 @@ LOG.addHandler(handler)
 
 
 class NovaConsumer(kombu.mixins.ConsumerMixin):
-    def __init__(self, name, connection, deployment, durable):
+    def __init__(self, name, connection, deployment, durable, queue_arguments):
         self.connection = connection
         self.deployment = deployment
         self.durable = durable
+        self.queue_arguments = queue_arguments
         self.name = name
         self.last_time = None
         self.pmi = None
         self.processed = 0
         self.total_processed = 0
 
+    def _create_exchange(self, name, type, exclusive=False, auto_delete=False):
+        return kombu.entity.Exchange(name, type=type, exclusive=exclusive,
+                                     durable=self.durable, auto_delete=auto_delete)
+
+    def _create_queue(self, name, nova_exchange, routing_key, exclusive=False,
+                     auto_delete=False):
+        return kombu.Queue(name, nova_exchange, durable=self.durable,
+                           auto_delete=exclusive, exclusive=auto_delete,
+                           queue_arguments=self.queue_arguments,
+                           routing_key=routing_key)
+
     def get_consumers(self, Consumer, channel):
-        nova_exchange = kombu.entity.Exchange("nova", type="topic",
+        nova_exchange = self._create_exchange("nova", "topic")
-                        exclusive=False, durable=self.durable,
-                        auto_delete=False)
 
         nova_queues = [
-            kombu.Queue("monitor.info", nova_exchange, durable=self.durable,
+            self._create_queue('monitor.info', nova_exchange, 'monitor.info'),
-                        auto_delete=False,
+            self._create_queue('monitor.error', nova_exchange, 'monitor.error')
-                        exclusive=False, routing_key='monitor.info'),
-            kombu.Queue("monitor.error", nova_exchange, durable=self.durable,
-                        auto_delete=False,
-                        exclusive=False, routing_key='monitor.error'),
         ]
 
         return [Consumer(queues=nova_queues, callbacks=[self.on_nova])]
 
-    def _process(self, body, message):
+    def _process(self, message):
         routing_key = message.delivery_info['routing_key']
-        payload = (routing_key, body)
-        jvalues = json.dumps(payload)
 
         body = str(message.body)
         args = (routing_key, json.loads(body))
@@ -110,12 +112,16 @@ class NovaConsumer(kombu.mixins.ConsumerMixin):
 
     def on_nova(self, body, message):
         try:
-            self._process(body, message)
+            self._process(message)
         except Exception, e:
             LOG.exception("Problem %s" % e)
         message.ack()
 
 
+def continue_running():
+    return True
+
+
 def run(deployment_config):
     name = deployment_config['name']
     host = deployment_config.get('rabbit_host', 'localhost')
@@ -124,8 +130,9 @@ def run(deployment_config):
     password = deployment_config.get('rabbit_password', 'rabbit')
     virtual_host = deployment_config.get('rabbit_virtual_host', '/')
     durable = deployment_config.get('durable_queue', True)
+    queue_arguments = deployment_config.get('queue_arguments', {})
 
-    deployment, new = models.get_or_create_deployment(name)
+    deployment, new = db.get_or_create_deployment(name)
 
     print "Starting worker for '%s'" % name
     LOG.info("%s: %s %s %s %s" % (name, host, port, user_id, virtual_host))
@@ -137,14 +144,16 @@ def run(deployment_config):
                   transport="librabbitmq",
                   virtual_host=virtual_host)
 
-    while True:
+    while continue_running():
         LOG.debug("Processing on '%s'" % name)
         with kombu.connection.BrokerConnection(**params) as conn:
             try:
-                consumer = NovaConsumer(name, conn, deployment, durable)
+                consumer = NovaConsumer(name, conn, deployment, durable,
+                                        queue_arguments)
                 consumer.run()
             except Exception as e:
                 LOG.exception("name=%s, exception=%s. Reconnecting in 5s" %
                                 (name, e))
                 time.sleep(5)
         LOG.debug("Completed processing on '%s'" % name)
+