AL-4081: Updated test for expected to fail case

Author: dubininsergey

labelbox/exceptions.py

@@ -129,3 +129,8 @@ class MALValidationError(LabelboxError):
 class OperationNotAllowedException(Exception):
     """Raised when user does not have permissions to a resource or has exceeded usage limit"""
     pass
+
+
+class ProcessingWaitTimeout(Exception):
+    """Raised when waiting for the data rows to be processed takes longer than allowed"""
+    pass

labelbox/schema/project.py

@@ -10,7 +10,7 @@
 import ndjson
 import requests
 from labelbox import utils
-from labelbox.exceptions import InvalidQueryError, LabelboxError
+from labelbox.exceptions import InvalidQueryError, LabelboxError, ProcessingWaitTimeout
 from labelbox.orm import query
 from labelbox.orm.db_object import DbObject, Deletable, Updateable
 from labelbox.orm.model import Entity, Field, Relationship
@@ -975,9 +975,9 @@ def _wait_until_data_rows_are_processed(self, data_row_ids: List[str], wait_proc
         start_time = datetime.now()
         while True:
             if (datetime.now() - start_time).total_seconds() >= wait_processing_max_seconds:
-                logger.warning(
-                    """Not all data rows have been processed, proceeding anyway""")
-                return
+                raise ProcessingWaitTimeout(
+                    "Maximum wait time exceeded while waiting for data rows to be processed. Try creating a batch a bit later"
+                )
 
             all_good = self.__check_data_rows_have_been_processed(data_row_ids)
             if all_good:

tests/integration/test_batch.py

@@ -1,5 +1,3 @@
-import warnings
-
 import pytest
 from labelbox import Dataset, Project
 from labelbox.schema.queue_mode import QueueMode
@@ -118,7 +116,7 @@ def test_batch_creation_with_processing_timeout(
     unique_dataset: Dataset
 ):
     """
-    Create a batch with zero wait time, this means that the waiting will termintate instantly
+    Create a batch with zero wait time, this means that the waiting logic will throw exception immediately
     """
     #  wait for these data rows to be processed
     valid_data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]
@@ -132,18 +130,13 @@ def test_batch_creation_with_processing_timeout(
         unique_dataset.export_data_rows())]
 
     data_row_ids = valid_data_rows + unprocessed_data_rows
-    with warnings.catch_warnings(record=True) as w:
-        warnings.simplefilter("always")
-        breakpoint()
+
+    with pytest.raises(ProcessingWaitTimeout):
         batch_project.create_batch(
             "batch to test failed data rows",
             data_row_ids,
             wait_processing_max_seconds=0
         )
-        assert len(w) == 1
-        assert issubclass(w[-1].category, DeprecationWarning)
-        assert "Not all data rows have been processed, proceeding anyway" in str(
-            w[-1].message)
 
 
 def test_export_data_rows(batch_project: Project, dataset: Dataset):