Export V1 to V2 for integration tests (#1618)

Author: Gabefire

libs/labelbox/src/labelbox/schema/export_task.py

@@ -920,3 +920,4 @@ def get_stream(
     def get_task(client, task_id):
         """Returns the task with the given id."""
         return ExportTask(Task.get_task(client, task_id))
+    

libs/labelbox/tests/conftest.py

@@ -929,7 +929,10 @@ def export_v2_test_helpers() -> Type[ExportV2Helpers]:
 
 @pytest.fixture
 def big_dataset_data_row_ids(big_dataset: Dataset):
-    yield [dr.uid for dr in list(big_dataset.export_data_rows())]
+    export_task = big_dataset.export()
+    export_task.wait_till_done()
+    stream = export_task.get_buffered_stream()
+    yield [dr.json["data_row"]["id"] for dr in stream]
 
 
 @pytest.fixture(scope='function')

libs/labelbox/tests/data/annotation_import/test_send_to_annotate_mea.py

@@ -55,7 +55,7 @@ def test_send_to_annotate_from_model(client, configured_project,
     # Check that the data row was sent to the new project
     destination_batches = list(destination_project.batches())
     assert len(destination_batches) == 1
-
+    
     destination_data_rows = list(destination_batches[0].export_data_rows())
     assert len(destination_data_rows) == len(data_row_ids)
     assert all([dr.uid in data_row_ids for dr in destination_data_rows])

libs/labelbox/tests/integration/conftest.py

@@ -314,7 +314,10 @@ def export_v2_test_helpers() -> Type[ExportV2Helpers]:
 
 @pytest.fixture
 def big_dataset_data_row_ids(big_dataset: Dataset):
-    yield [dr.uid for dr in list(big_dataset.export_data_rows())]
+    export_task = big_dataset.export()
+    export_task.wait_till_done()
+    stream = export_task.get_buffered_stream()
+    yield [dr.json["data_row"]["id"] for dr in stream]
 
 
 @pytest.fixture(scope='function')

libs/labelbox/tests/integration/test_batch.py

@@ -1,7 +1,6 @@
 import time
 from typing import List
 from uuid import uuid4
-
 import pytest
 
 from labelbox import Dataset, Project
@@ -92,7 +91,10 @@ def test_create_batch_async(project: Project,
 
 def test_create_batch_with_consensus_settings(project: Project,
                                               small_dataset: Dataset):
-    data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]
+    export_task = small_dataset.export()
+    export_task.wait_till_done()
+    stream = export_task.get_buffered_stream()
+    data_rows = [dr.json["data_row"]["id"] for dr in stream]
     consensus_settings = {"coverage_percentage": 0.1, "number_of_labels": 3}
     batch = project.create_batch("batch with consensus settings",
                                  data_rows,
@@ -105,31 +107,44 @@ def test_create_batch_with_consensus_settings(project: Project,
 
 def test_create_batch_with_data_row_class(project: Project,
                                           small_dataset: Dataset):
-    data_rows = list(small_dataset.export_data_rows())
+    export_task = small_dataset.export()
+    export_task.wait_till_done()
+    stream = export_task.get_buffered_stream()
+    data_rows = [dr.json["data_row"]["id"] for dr in stream]
     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(project: Project, small_dataset: Dataset):
-    data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]
+    export_task = small_dataset.export()
+    export_task.wait_till_done()
+    stream = export_task.get_buffered_stream()
+    data_rows = [dr.json["data_row"]["id"] for dr in stream]
+    
     batch = project.create_batch("batch to archive", data_rows)
     batch.remove_queued_data_rows()
-    exported_data_rows = list(batch.export_data_rows())
-
-    assert len(exported_data_rows) == 0
+    overview = project.get_overview()
+    
+    assert overview.to_label == 0
 
 
 def test_delete(project: Project, small_dataset: Dataset):
-    data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]
+    export_task = small_dataset.export()
+    export_task.wait_till_done()
+    stream = export_task.get_buffered_stream()
+    data_rows = [dr.json["data_row"]["id"] for dr in stream]
     batch = project.create_batch("batch to delete", data_rows)
     batch.delete()
 
     assert len(list(project.batches())) == 0
 
 
 def test_batch_project(project: Project, small_dataset: Dataset):
-    data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]
+    export_task = small_dataset.export()
+    export_task.wait_till_done()
+    stream = export_task.get_buffered_stream()
+    data_rows = [dr.json["data_row"]["id"] for dr in stream]
     batch = project.create_batch("batch to test project relationship",
                                  data_rows)
 
@@ -186,6 +201,7 @@ def test_batch_creation_with_processing_timeout(
     project._wait_processing_max_seconds = stashed_wait_timeout
 
 
+@pytest.mark.export_v1("export_v1 test remove later")
 def test_export_data_rows(project: Project, dataset: Dataset, image_url: str,
                           external_id: str):
     n_data_rows = 2
@@ -255,18 +271,34 @@ def test_list_project_batches_with_no_batches(project: Project):
     reason="Test cannot be used effectively with MAL/LabelImport. \
 Fix/Unskip after resolving deletion with MAL/LabelImport")
 def test_delete_labels(project, small_dataset):
-    data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]
+    export_task = small_dataset.export()
+    export_task.wait_till_done()
+    stream = export_task.get_buffered_stream()
+    data_rows = [dr.json["data_row"]["id"] for dr in stream]
     batch = project.create_batch("batch to delete labels", data_rows)
 
 
 @pytest.mark.skip(
     reason="Test cannot be used effectively with MAL/LabelImport. \
 Fix/Unskip after resolving deletion with MAL/LabelImport")
 def test_delete_labels_with_templates(project: Project, small_dataset: Dataset):
-    data_rows = [dr.uid for dr in list(small_dataset.export_data_rows())]
+    export_task = small_dataset.export()
+    export_task.wait_till_done()
+    stream = export_task.get_buffered_stream()
+    data_rows = [dr.json["data_row"]["id"] for dr in stream]
     batch = project.create_batch("batch to delete labels w templates",
                                  data_rows)
-    exported_data_rows = list(batch.export_data_rows())
+
+    export_task = project.export(filters={"batch_ids": [batch.uid]})
+    export_task.wait_till_done()
+    stream = export_task.get_buffered_stream()
+    exported_data_rows = [dr.json["data_row"]["id"] for dr in stream]
+
     res = batch.delete_labels(labels_as_template=True)
-    exported_data_rows = list(batch.export_data_rows())
+
+    export_task = project.export(filters={"batch_ids": [batch.uid]})
+    export_task.wait_till_done()
+    stream = export_task.get_buffered_stream()
+    exported_data_rows = [dr.json["data_row"]["id"] for dr in stream]
+
     assert len(exported_data_rows) == 5

libs/labelbox/tests/integration/test_batches.py

@@ -20,7 +20,10 @@ def test_create_batches(project: Project, big_dataset_data_row_ids: List[str]):
 
 
 def test_create_batches_from_dataset(project: Project, big_dataset: Dataset):
-    data_rows = [dr.uid for dr in list(big_dataset.export_data_rows())]
+    export_task = big_dataset.export()
+    export_task.wait_till_done()
+    stream = export_task.get_buffered_stream()
+    data_rows = [dr.json["data_row"]["id"] for dr in stream]
     project._wait_until_data_rows_are_processed(data_rows, [], 300)
 
     task = project.create_batches_from_dataset("test-batch",

libs/labelbox/tests/integration/test_project.py

@@ -1,7 +1,6 @@
 import time
 import os
 import uuid
-
 import pytest
 import requests
 
@@ -201,8 +200,10 @@ def test_batches(project: Project, dataset: Dataset, image_url):
         },
     ] * 2)
     task.wait_till_done()
-    # TODO: Move to export_v2
-    data_rows = [dr.uid for dr in list(dataset.export_data_rows())]
+    export_task = dataset.export()
+    export_task.wait_till_done()
+    stream = export_task.get_buffered_stream()
+    data_rows = [dr.json["data_row"]["id"] for dr in stream]
     batch_one = f'batch one {uuid.uuid4()}'
     batch_two = f'batch two {uuid.uuid4()}'
     project.create_batch(batch_one, [data_rows[0]])
@@ -217,19 +218,17 @@ def test_create_batch_with_global_keys_sync(project: Project, data_rows):
     global_keys = [dr.global_key for dr in data_rows]
     batch_name = f'batch {uuid.uuid4()}'
     batch = project.create_batch(batch_name, global_keys=global_keys)
-    # TODO: Move to export_v2
-    batch_data_rows = set(batch.export_data_rows())
-    assert batch_data_rows == set(data_rows)
+    
+    assert batch.size == len(set(data_rows))
 
 
 @pytest.mark.parametrize('data_rows', [2], indirect=True)
 def test_create_batch_with_global_keys_async(project: Project, data_rows):
     global_keys = [dr.global_key for dr in data_rows]
     batch_name = f'batch {uuid.uuid4()}'
     batch = project._create_batch_async(batch_name, global_keys=global_keys)
-    # TODO: Move to export_v2
-    batch_data_rows = set(batch.export_data_rows())
-    assert batch_data_rows == set(data_rows)
+    
+    assert batch.size == len(set(data_rows))
 
 
 def test_media_type(client, project: Project, rand_gen):

libs/labelbox/tests/integration/test_send_to_annotate.py

@@ -1,11 +1,12 @@
-from labelbox import UniqueIds, OntologyBuilder, LabelingFrontend
+from labelbox import UniqueIds, Project, Ontology, Client
 from labelbox.schema.conflict_resolution_strategy import ConflictResolutionStrategy
+from typing import List
 
 
 def test_send_to_annotate_include_annotations(
-        client, configured_batch_project_with_label, project_pack, ontology):
+        client: Client, configured_batch_project_with_label: Project, project_pack: List[Project], ontology: Ontology):
     [source_project, _, data_row, _] = configured_batch_project_with_label
-    destination_project = project_pack[0]
+    destination_project: Project = project_pack[0]
 
     src_ontology = source_project.ontology()
     destination_project.setup_editor(ontology)
@@ -45,10 +46,14 @@ def test_send_to_annotate_include_annotations(
         # Check that the data row was sent to the new project
         destination_batches = list(destination_project.batches())
         assert len(destination_batches) == 1
-
-        destination_data_rows = list(destination_batches[0].export_data_rows())
+        
+        export_task = destination_project.export()
+        export_task.wait_till_done()
+        stream = export_task.get_buffered_stream()
+        
+        destination_data_rows = [dr.json["data_row"]["id"] for dr in stream]
         assert len(destination_data_rows) == 1
-        assert destination_data_rows[0].uid == data_row.uid
+        assert destination_data_rows[0] == data_row.uid
 
         # Verify annotations were copied into the destination project
         destination_project_labels = (list(destination_project.labels()))

libs/labelbox/tests/integration/test_task.py

@@ -58,10 +58,9 @@ def test_task_success_json(dataset, image_url, snapshot):
                           'test_task.test_task_success_json.json')
     assert len(task.result)
 
-
+@pytest.mark.export_v1("export_v1 test remove later")
 def test_task_success_label_export(client, configured_project_with_label):
     project, _, _, _ = configured_project_with_label
-    # TODO: Move to export_v2
     project.export_labels()
     user = client.get_user()
     task = None