Merge pull request #103 from gwu-libraries/t91-job_improvements

T91 job improvements

README.md

@@ -10,3 +10,8 @@ sfm-ui provides a Django app which:
 - Publishes harvest.start messages for flickr SeedSets.  The app schedules harvest.start messages for publication when the user updates an existing, active SeedSet.
 - Includes a scheduler which uses [apscheduler](http://apscheduler.readthedocs.org) to schedule publication of harvest.start messages.
 - Binds to `harvest.status.*(.*)` messages and creates a Harvest object (visible in the admin views) for each harvest status message received.  The message consumer is started via the `startconsumer` management command.
+
+## Unit tests
+
+     cd sfm
+    ./manage.py test --settings=sfm.settings.test_settings

sfm/sfm/settings/common.py

@@ -90,7 +90,7 @@
         'ENGINE': 'django.db.backends.postgresql_psycopg2',
         'NAME': 'sfmdatabase',
         'USER': 'postgres',
-        'PASSWORD': env['DB_ENV_POSTGRES_PASSWORD'],
+        'PASSWORD': env.get('DB_ENV_POSTGRES_PASSWORD'),
         'HOST': 'db',
         'PORT': '5432',
     }

sfm/sfm/settings/test_settings.py

@@ -0,0 +1,9 @@
+from .common import *
+
+DATABASES = {
+    # for unit tests
+    'default': {
+        'ENGINE': 'django.db.backends.sqlite3',
+        'NAME': 'testdb'
+    }
+}

sfm/ui/jobs.py

@@ -1,71 +1,60 @@
 import json
 from .rabbit import RabbitWorker
-from .models import SeedSet, Collection, Seed, Credential
+from .models import SeedSet
+from django.core.exceptions import ObjectDoesNotExist
 import logging
+import datetime
 
 log = logging.getLogger(__name__)
 
 
-def seedset_harvest(d):
-    # To get value of Collection id for the associated collection object.
-    for collection_id in list(Collection.objects.filter(
-            id=SeedSet.objects.filter(id=d).values('collection')).values('id')):
-        if 'id' in collection_id:
-            value = collection_id['id']
+def seedset_harvest(seedset_id):
 
-    # To get value of the token for the associated credential object.
-    for token in list(Credential.objects.filter(
-        id=SeedSet.objects.filter(id=d).values('credential')).values(
-            'token')):
-        if 'token' in token:
-            credential = token['token']
+    message = {
+        "collection": {},
+        "seeds": []
+    }
 
-    # To get value of platform
-    for platform in list(Credential.objects.filter(
-        id=SeedSet.objects.filter(id=d).values('credential')).values(
-            'platform')):
-        if 'platform' in platform:
-            media = platform['platform']
+    # Retrieve seedset
+    try:
+        seedset = SeedSet.objects.get(id=seedset_id)
+    except ObjectDoesNotExist:
+        log.error("Harvesting seedset %s failed because seedset does not exist", seedset_id)
+        return
 
-    # To get list of seeds
-    seeds = list(Seed.objects.filter(seed_set=d).select_related(
-        'seeds').values('token', 'uid'))
-    # To remove empty token values from the list of seeds --
-    # Need to update below code
-    #
-    # if item['token'] not in seeds:
-        # item.pop('token', None)
-    for item in seeds:
-        if item['token'] == '':
-            item.pop('token', None)
+    # Id
+    harvest_id = "harvest:{}:{}".format(seedset_id, datetime.datetime.now().isoformat())
+    message["id"] = harvest_id
 
-    # To get harvest type, options and credentials
-    harvest_type = SeedSet.objects.filter(id=d).values(
-        'harvest_type')[0]["harvest_type"]
-    options = json.loads(SeedSet.objects.filter(id=d).values(
-        'harvest_options')[0]["harvest_options"])
-    credential = json.loads(str(credential))
+    # Collection
+    collection = seedset.collection
+    message["collection"]["id"] = "collection:{}".format(collection.id)
+    # TODO: Different approach for path
+    message["collection"]["path"] = "/tmp/collection/{}".format(collection.id)
 
-    # Routing Key
-    key = ''.join(['harvest.start.',str(media),'.',harvest_type])
+    # Credential
+    credentials = seedset.credential
+    message["credentials"] = json.loads(str(credentials.token))
 
-    # message to be sent to queue
-    # TODO: Unique id
-    # TODO: Correct path
-    m = {
-        'id': str(d),
-        'type': harvest_type,
-        'options': options,
-        'credentials': credential,
-        'collection': {
-            'id': str(value),
-            'path': '/tmp/collection/'+str(value)
-        },
-        'seeds': seeds,
-    }
+    # Type
+    harvest_type = seedset.harvest_type
+    message["type"] = harvest_type
+
+    # Options
+    message["options"] = json.loads(seedset.harvest_options or "{}")
+
+    # Seeds
+    for seed in seedset.seeds.all():
+        seed_map = {}
+        if seed.token:
+            seed_map["token"] = seed.token
+        if seed.uid:
+            seed_map["uid"] = seed.uid
+        message["seeds"].append(seed_map)
+
+    routing_key = "harvest.start.{}.{}".format(credentials.platform, harvest_type)
 
-    log.debug("Sending %s message to %s with id %s", harvest_type, key,
-              m["id"])
+    log.debug("Sending %s message to %s with id %s", harvest_type, routing_key, harvest_id)
 
     # Publish message to queue via rabbit worker
-    RabbitWorker().send_message(m, key)
+    RabbitWorker().send_message(message, routing_key)

sfm/ui/rabbit.py

@@ -30,9 +30,12 @@ def get_connection():
                           password=settings.RABBITMQ_PASSWORD)
 
     def declare_exchange(self):
-        with self.get_connection() as connection:
-            log.debug("Declaring %s exchange", self.exchange.name)
-            self.exchange(connection).declare()
+        try:
+            with self.get_connection() as connection:
+                log.debug("Declaring %s exchange", self.exchange.name)
+                self.exchange(connection).declare()
+        except:
+            log.error("Error connecting to RabbitMQ to declare exchange")
 
     def send_message(self, message, routing_key):
         with self.get_connection() as connection:

sfm/ui/test_jobs.py

@@ -0,0 +1,48 @@
+from django.test import TestCase
+import json
+from mock import MagicMock, patch
+from .jobs import seedset_harvest
+from .models import SeedSet, Collection, Seed, Credential, Group, User
+from .rabbit import RabbitWorker
+
+
+class JobsTests(TestCase):
+    def setUp(self):
+        self.user = User.objects.create_superuser(username="test_user", email="test_user@test.com",
+                                                  password="test_password")
+        self.group = Group.objects.create(name="test_group")
+        self.collection = Collection.objects.create(group=self.group, name="test_collection")
+        self.credential_token = {"key": "test_key"}
+        self.credential = Credential.objects.create(user=self.user, platform="test_platform",
+                                                    token=json.dumps(self.credential_token))
+        self.harvest_options = {"test_option": "test_value"}
+        self.seedset = SeedSet.objects.create(collection=self.collection, credential=self.credential,
+                                              harvest_type="test_type", name="test_seedset",
+                                              harvest_options=json.dumps(self.harvest_options))
+        Seed.objects.create(seed_set=self.seedset, token="test_token1")
+        Seed.objects.create(seed_set=self.seedset, uid="test_uid2")
+        Seed.objects.create(seed_set=self.seedset, token="test_token3", uid="test_uid3")
+
+    @patch("ui.jobs.RabbitWorker", autospec=True)
+    def test_seedset_harvest(self, mock_rabbit_worker_class):
+        mock_rabbit_worker = MagicMock(spec=RabbitWorker)
+        mock_rabbit_worker_class.side_effect = [mock_rabbit_worker]
+
+        seedset_harvest(self.seedset.id)
+
+        name, args, kwargs = mock_rabbit_worker.mock_calls[0]
+        self.assertEqual("send_message", name)
+        message = args[0]
+        self.assertEqual("collection:{}".format(self.collection.id), message["collection"]["id"])
+        self.assertEqual("/tmp/collection/{}".format(self.collection.id), message["collection"]["path"])
+        self.assertDictEqual(self.harvest_options, message["options"])
+        self.assertDictEqual({"token": "test_token1"}, message["seeds"][0])
+        self.assertDictEqual({"uid": "test_uid2"}, message["seeds"][1])
+        self.assertDictEqual({"token": "test_token3", "uid": "test_uid3"}, message["seeds"][2])
+        self.assertEqual("test_type", message["type"])
+        self.assertTrue(message["id"].startswith("harvest:{}:".format(self.seedset.id)))
+        self.assertEqual("harvest.start.test_platform.test_type", args[1])
+
+    def test_missing_seedset_harvest(self):
+        # Error should be logged and nothing happens
+        seedset_harvest(1234567)