Fix unit tests

arkindex/dataimport/tasks.py

@@ -124,7 +124,7 @@ def save_ml_results(self, results, **kwargs):
     pages = Page.objects.filter(pk__in=list(results.keys()))
     for page in pages:
         result = results.get(page.id)
-        if result is not None:
+        if result is None:
             continue
 
         page.classification = result['classification']

arkindex/dataimport/tests/test_tasks.py

 from arkindex.project.tests import RedisMockAPITestCase
-from arkindex.dataimport.tasks import save_classification
+from arkindex.dataimport.tasks import save_ml_results
 from arkindex.documents.models import Page, Corpus
 
 
@@ -7,35 +7,43 @@ class TestTasks(RedisMockAPITestCase):
     """
     Test data imports tasks
     """
-    def test_save_classification(self):
+    def test_save_ml_results(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,
-                }
-            ]
+            dog.id: {
+                'classification': [
+                    {
+                        'label': 'dog',
+                        'probability': 0.9,
+                    }
+                ],
+                'text': 'This is a dog',
+            },
+            cat.id: {
+                'classification': [
+                    {
+                        'label': 'cat',
+                        'probability': 0.8,
+                    }
+                ],
+                'text': 'This is a cat. meow',
+            },
         }
-        save_classification(classification)
+        save_ml_results(classification)
 
         dog.refresh_from_db()
         self.assertEqual(dog.classification, [{
             'label': 'dog',
             'probability': 0.9,
         }])
+        self.assertEqual(dog.text, 'This is a dog')
 
         cat.refresh_from_db()
         self.assertEqual(cat.classification, [{
             'label': 'cat',
             'probability': 0.8,
         }])
+        self.assertEqual(cat.text, 'This is a cat. meow')