diff --git a/.coverage b/.coverage deleted file mode 100644 index 9b42c3f..0000000 Binary files a/.coverage and /dev/null differ diff --git a/.gitignore b/.gitignore index 54ef793..618391f 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,5 @@ *.pyc local_settings.py +.coverage +worker.log diff --git a/run_tests.sh b/run_tests.sh index 54904d1..d5e661f 100755 --- a/run_tests.sh +++ b/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 diff --git a/stacktach/db.py b/stacktach/db.py index a00d70e..23db4fb 100644 --- a/stacktach/db.py +++ b/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) diff --git a/stacktach/models.py b/stacktach/models.py index 3ae2ae1..ca2ea94 100644 --- a/stacktach/models.py +++ b/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, diff --git a/tests/unit/test_worker.py b/tests/unit/test_worker.py new file mode 100644 index 0000000..ada86bb --- /dev/null +++ b/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() \ No newline at end of file diff --git a/worker/__init__.py b/worker/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/worker/worker.py b/worker/worker.py index 9a3165f..db7a22b 100644 --- a/worker/worker.py +++ b/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 datetime_to_decimal as dt +from stacktach import db, views 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", - exclusive=False, durable=self.durable, - auto_delete=False) + nova_exchange = self._create_exchange("nova", "topic") nova_queues = [ - kombu.Queue("monitor.info", nova_exchange, durable=self.durable, - auto_delete=False, - exclusive=False, routing_key='monitor.info'), - kombu.Queue("monitor.error", nova_exchange, durable=self.durable, - auto_delete=False, - exclusive=False, routing_key='monitor.error'), + self._create_queue('monitor.info', nova_exchange, 'monitor.info'), + self._create_queue('monitor.error', nova_exchange, '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) +