Reduce wait time for DR processing in tests (#1147)

Author: kkim-labelbox

.github/workflows/python-package.yml

@@ -89,4 +89,4 @@ jobs:
 
           DA_GCP_LABELBOX_API_KEY: ${{ secrets[matrix.da-test-key] }}
         run: |
-          tox -e py -- -n 10 -svv --reruns 5 --reruns-delay 10
+          tox -e py -- -n 10 -svv --reruns 3 --reruns-delay 3

.pre-commit-config.yaml

@@ -1,6 +1,6 @@
 repos:
   - repo: https://github.com/google/yapf
-    rev: v0.44.0
+    rev: v0.40.1
     hooks:
       - id: yapf
         name: "yapf"

tests/integration/annotation_import/conftest.py

@@ -11,6 +11,9 @@
 from labelbox.schema.annotation_import import LabelImport, AnnotationImportState
 from labelbox.schema.queue_mode import QueueMode
 
+DATA_ROW_PROCESSING_WAIT_TIMEOUT_SECONDS = 30
+DATA_ROW_PROCESSING_WAIT_SLEEP_INTERNAL_SECONDS = 5
+
 
 @pytest.fixture()
 def audio_data_row(rand_gen):
@@ -486,7 +489,10 @@ def configured_project(client, ontology, rand_gen, image_url):
 
     for _ in range(len(ontology['tools']) + len(ontology['classifications'])):
         data_row_ids.append(dataset.create_data_row(row_data=image_url).uid)
-    project._wait_until_data_rows_are_processed(data_row_ids=data_row_ids)
+    project._wait_until_data_rows_are_processed(
+        data_row_ids=data_row_ids,
+        wait_processing_max_seconds=DATA_ROW_PROCESSING_WAIT_TIMEOUT_SECONDS,
+        sleep_interval=DATA_ROW_PROCESSING_WAIT_SLEEP_INTERNAL_SECONDS)
     project.datasets.connect(dataset)
     project.data_row_ids = data_row_ids
     yield project
@@ -505,7 +511,10 @@ def configured_project_pdf(client, ontology, rand_gen, pdf_url):
     project.setup(editor, ontology)
     data_row_ids = []
     data_row_ids.append(dataset.create_data_row(pdf_url).uid)
-    project._wait_until_data_rows_are_processed(data_row_ids=data_row_ids)
+    project._wait_until_data_rows_are_processed(
+        data_row_ids=data_row_ids,
+        wait_processing_max_seconds=DATA_ROW_PROCESSING_WAIT_TIMEOUT_SECONDS,
+        sleep_interval=DATA_ROW_PROCESSING_WAIT_SLEEP_INTERNAL_SECONDS)
     project.datasets.connect(dataset)
     project.data_row_ids = data_row_ids
     yield project

tests/integration/conftest.py

@@ -22,6 +22,8 @@
 from labelbox.schema.user import User
 
 IMG_URL = "https://picsum.photos/200/300.jpg"
+DATA_ROW_PROCESSING_WAIT_TIMEOUT_SECONDS = 30
+DATA_ROW_PROCESSING_WAIT_SLEEP_INTERNAL_SECONDS = 5
 
 
 class Environ(Enum):
@@ -398,7 +400,10 @@ def configured_batch_project_with_label(client, rand_gen, image_url,
     One label is already created and yielded when using fixture
     """
     data_rows = [dr.uid for dr in list(dataset.data_rows())]
-    batch_project._wait_until_data_rows_are_processed(data_row_ids=data_rows)
+    batch_project._wait_until_data_rows_are_processed(
+        data_row_ids=data_rows,
+        wait_processing_max_seconds=DATA_ROW_PROCESSING_WAIT_TIMEOUT_SECONDS,
+        sleep_interval=DATA_ROW_PROCESSING_WAIT_SLEEP_INTERNAL_SECONDS)
     batch_project.create_batch("test-batch", data_rows)
 
     ontology = _setup_ontology(batch_project)
@@ -421,7 +426,10 @@ def configured_batch_project_with_multiple_datarows(batch_project, dataset,
     """
     global_keys = [dr.global_key for dr in data_rows]
 
-    batch_project._wait_until_data_rows_are_processed(global_keys=global_keys)
+    batch_project._wait_until_data_rows_are_processed(
+        global_keys=global_keys,
+        wait_processing_max_seconds=DATA_ROW_PROCESSING_WAIT_TIMEOUT_SECONDS,
+        sleep_interval=DATA_ROW_PROCESSING_WAIT_SLEEP_INTERNAL_SECONDS)
     batch_name = f'batch {uuid.uuid4()}'
     batch_project.create_batch(batch_name, global_keys=global_keys)
 

tests/integration/test_batch.py

@@ -57,8 +57,6 @@ def test_create_batch(batch_project: Project, big_dataset: Dataset):
 
 def test_create_batch_async(batch_project: Project, big_dataset: Dataset):
     data_rows = [dr.uid for dr in list(big_dataset.export_data_rows())]
-    batch_project._wait_until_data_rows_are_processed(
-        data_rows, batch_project._wait_processing_max_seconds)
     batch = batch_project._create_batch_async("big-batch",
                                               data_rows,
                                               priority=3)
@@ -146,8 +144,6 @@ def test_batch_creation_with_processing_timeout(batch_project: Project,
     """
     #  wait for these data rows to be processed
     valid_data_rows = [dr.uid for dr in list(small_dataset.data_rows())]
-    batch_project._wait_until_data_rows_are_processed(
-        valid_data_rows, wait_processing_max_seconds=3600, sleep_interval=5)
 
     # upload data rows for this dataset and don't wait
     upload_invalid_data_rows_for_dataset(unique_dataset)