add project.batches() test

Author: farkob

labelbox/pagination.py

@@ -25,7 +25,7 @@ def __init__(self,
                  params: Dict[str, str],
                  dereferencing: Union[List[str], Dict[str, Any]],
                  obj_class: Union[Type["DbObject"], Callable[[Any, Any], Any]],
-                 cursor_path: Optional[Dict[str, Any]] = None,
+                 cursor_path: Optional[List[str]] = None,
                  experimental: bool = False):
         """ Creates a PaginatedCollection.
 
@@ -105,7 +105,7 @@ def get_next_page(self) -> Tuple[Dict[str, Any], bool]:
 
 class _CursorPagination(_Pagination):
 
-    def __init__(self, cursor_path: Dict[str, Any], *args, **kwargs):
+    def __init__(self, cursor_path: List[str], *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.cursor_path = cursor_path
         self.next_cursor: Optional[Any] = None

labelbox/schema/project.py

@@ -855,12 +855,7 @@ def batches(self) -> PaginatedCollection:
             self.client,
             query_str, {id_param: self.uid}, ['project', 'batches', 'nodes'],
             lambda client, res: Entity.Batch(client, self.uid, res),
-            cursor_path={
-                'project': None,
-                'batches': None,
-                'pageInfo': None,
-                'endCursor': None
-            },
+            cursor_path=['project', 'batches', 'pageInfo', 'endCursor'],
             experimental=True)
 
     def upload_annotations(

tests/integration/test_project.py

@@ -4,7 +4,7 @@
 import pytest
 import requests
 
-from labelbox import Project, LabelingFrontend
+from labelbox import Project, LabelingFrontend, Dataset
 from labelbox.exceptions import InvalidQueryError
 
 
@@ -201,3 +201,22 @@ def test_queue_mode(configured_project: Project):
     ) == configured_project.QueueMode.Dataset
     configured_project.update(queue_mode=configured_project.QueueMode.Batch)
     assert configured_project.queue_mode() == configured_project.QueueMode.Batch
+
+
+def test_batches(configured_project: Project, dataset: Dataset, image_url):
+    task = dataset.create_data_rows([
+        {
+            "row_data": image_url,
+            "external_id": "my-image"
+        },
+    ] * 2)
+    task.wait_till_done()
+    configured_project.update(queue_mode=configured_project.QueueMode.Batch)
+    data_rows = [dr.uid for dr in list(dataset.export_data_rows())]
+    batch_one = 'batch one'
+    batch_two = 'batch two'
+    configured_project.create_batch(batch_one, [data_rows[0]])
+    configured_project.create_batch(batch_two, [data_rows[1]])
+
+    names = set([batch.name for batch in list(configured_project.batches())])
+    assert names == set([batch_one, batch_two])