Add test

arkindex/dataimport/tasks.py

@@ -45,7 +45,6 @@ def check_images(self, dataimport):
         fmt = None
 
         try:
-            print(datafile.staging_path)
             fmt = Image.open(datafile.staging_path).format
             assert fmt in ('JPEG', 'JPEG2000')
         except IOError as e:

arkindex/dataimport/tests/test_tasks.py

+from arkindex.project.tests import RedisMockAPITestCase
+from arkindex.dataimport.tasks import save_classification
+from arkindex.documents.models import Page, Corpus
+
+
+class TestTasks(RedisMockAPITestCase):
+    """
+    Test data imports tasks
+    """
+    def test_save_classification(self):
+        corpus = Corpus.objects.create(name='test class')
+        dog = Page.objects.create(corpus=corpus, name='A dog')
+        cat = Page.objects.create(corpus=corpus, name='A cat')
+
+        classification = {
+            dog.id: [
+                {
+                    'label': 'dog',
+                    'probability': 0.9,
+                }
+            ],
+            cat.id: [
+                {
+                    'label': 'cat',
+                    'probability': 0.8,
+                }
+            ]
+        }
+        save_classification(classification)
+
+        dog.refresh_from_db()
+        self.assertEqual(dog.classification, [{
+            'label': 'dog',
+            'probability': 0.9,
+        }])
+
+        cat.refresh_from_db()
+        self.assertEqual(cat.classification, [{
+            'label': 'cat',
+            'probability': 0.8,
+        }])

arkindex/project/tests/__init__.py

@@ -50,6 +50,10 @@ class RedisMockMixin(object):
         for m in self.mocked:
             m.return_value = self.redis
 
+        # Patch the add_message
+        self.messages = patch('arkindex.project.celery.ExtendedRedisBackend.add_message')
+        self.messages.start()
+
     def tearDown(self):
         for p in self.patches:
             p.stop()