WIP

tests/web/test_admin_routes_enhanced.py

@@ -31,6 +31,7 @@ class MockAdminDocument:
         self.batch_id = kwargs.get('batch_id', None)
         self.csv_field_values = kwargs.get('csv_field_values', None)
         self.annotation_lock_until = kwargs.get('annotation_lock_until', None)
+        self.category = kwargs.get('category', 'invoice')
         self.created_at = kwargs.get('created_at', datetime.utcnow())
         self.updated_at = kwargs.get('updated_at', datetime.utcnow())
 
@@ -67,12 +68,13 @@ class MockAdminDB:
 
     def get_documents_by_token(
         self,
-        admin_token,
+        admin_token=None,
         status=None,
         upload_source=None,
         has_annotations=None,
         auto_label_status=None,
         batch_id=None,
+        category=None,
         limit=20,
         offset=0
     ):
@@ -95,6 +97,8 @@ class MockAdminDB:
             docs = [d for d in docs if d.auto_label_status == auto_label_status]
         if batch_id:
             docs = [d for d in docs if str(d.batch_id) == str(batch_id)]
+        if category:
+            docs = [d for d in docs if d.category == category]
 
         total = len(docs)
         return docs[offset:offset+limit], total