[X-0] Fix create batch with data row classes

Author: mkozik-lb

labelbox/schema/project.py

@@ -606,7 +606,7 @@ def create_batch(self,
             raise ValueError("You need at least one data row in a batch")
 
         self._wait_until_data_rows_are_processed(
-            data_rows, self._wait_processing_max_seconds)
+            dr_ids, self._wait_processing_max_seconds)
         method = 'createBatchV2'
         query_str = """mutation %sPyApi($projectId: ID!, $batchInput: CreateBatchInput!) {
               project(where: {id: $projectId}) {

tests/integration/test_batch.py

@@ -56,8 +56,8 @@ def test_create_batch(batch_project: Project, big_dataset: Dataset):
 
 
 def test_create_batch_with_consensus_settings(batch_project: Project,
-                                              big_dataset: Dataset):
-    data_rows = [dr.uid for dr in list(big_dataset.export_data_rows())]
+                                              small_dataset: Dataset):
+    data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]
     consensus_settings = {"coverage_percentage": 0.1, "number_of_labels": 3}
     batch = batch_project.create_batch("batch with consensus settings",
                                        data_rows,
@@ -68,6 +68,14 @@ def test_create_batch_with_consensus_settings(batch_project: Project,
     assert batch.consensus_settings == consensus_settings
 
 
+def test_create_batch_with_data_row_class(batch_project: Project,
+                                          small_dataset: Dataset):
+    data_rows = list(small_dataset.export_data_rows())
+    batch = batch_project.create_batch("test-batch-data-rows", data_rows, 3)
+    assert batch.name == "test-batch-data-rows"
+    assert batch.size == len(data_rows)
+
+
 def test_archive_batch(batch_project: Project, small_dataset: Dataset):
     data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]
     batch = batch_project.create_batch("batch to archive", data_rows)