[SDK-428] Add integration tests for ExportTask (#1300)

Author: kopreschko

tests/integration/export_v2/conftest.py tests/integration/export/conftest.pytests/integration/export_v2/conftest.py renamed to tests/integration/export/conftest.py

@@ -314,6 +314,8 @@ def model_run_with_data_rows(client, configured_project_with_ontology,
                              model_run_predictions, model_run,
                              wait_for_label_processing):
     configured_project_with_ontology.enable_model_assisted_labeling()
+    use_data_row_ids = [p['dataRow']['id'] for p in model_run_predictions]
+    model_run.upsert_data_rows(use_data_row_ids)
 
     upload_task = LabelImport.create_from_objects(
         client, configured_project_with_ontology.uid,
@@ -326,7 +328,7 @@ def model_run_with_data_rows(client, configured_project_with_ontology,
     labels = wait_for_label_processing(configured_project_with_ontology)
     label_ids = [label.uid for label in labels]
     model_run.upsert_labels(label_ids)
-    yield model_run
+    yield model_run, labels
     model_run.delete()
     # TODO: Delete resources when that is possible ..
 

…ation/export_v2/test_export_data_rows.py …n/export/legacy/test_export_data_rows.pytests/integration/export_v2/test_export_data_rows.py renamed to tests/integration/export/legacy/test_export_data_rows.py tests/integration/export_v2/test_export_data_rows.py renamed to tests/integration/export/legacy/test_export_data_rows.py

@@ -1,37 +1,29 @@
 import time
 from labelbox import DataRow
-from labelbox.schema.media_type import MediaType
 
 
 def test_export_data_rows(client, data_row, wait_for_data_row_processing):
     # Ensure created data rows are indexed
     data_row = wait_for_data_row_processing(client, data_row)
     time.sleep(7)  # temp fix for ES indexing delay
-    params = {
-        "include_performance_details": True,
-        "include_labels": True,
-        "media_type_override": MediaType.Image,
-        "project_details": True,
-        "data_row_details": True
-    }
 
     task = DataRow.export_v2(client=client, data_rows=[data_row])
     task.wait_till_done()
     assert task.status == "COMPLETE"
     assert task.errors is None
     assert len(task.result) == 1
-    assert task.result[0]['data_row']['id'] == data_row.uid
+    assert task.result[0]["data_row"]["id"] == data_row.uid
 
     task = DataRow.export_v2(client=client, data_rows=[data_row.uid])
     task.wait_till_done()
     assert task.status == "COMPLETE"
     assert task.errors is None
     assert len(task.result) == 1
-    assert task.result[0]['data_row']['id'] == data_row.uid
+    assert task.result[0]["data_row"]["id"] == data_row.uid
 
     task = DataRow.export_v2(client=client, global_keys=[data_row.global_key])
     task.wait_till_done()
     assert task.status == "COMPLETE"
     assert task.errors is None
     assert len(task.result) == 1
-    assert task.result[0]['data_row']['id'] == data_row.uid
+    assert task.result[0]["data_row"]["id"] == data_row.uid

…gration/export_v2/test_export_dataset.py …ion/export/legacy/test_export_dataset.pytests/integration/export_v2/test_export_dataset.py renamed to tests/integration/export/legacy/test_export_dataset.py tests/integration/export_v2/test_export_dataset.py renamed to tests/integration/export/legacy/test_export_dataset.py

@@ -0,0 +1,45 @@
+import time
+
+
+def _model_run_export_v2_results(model_run, task_name, params, num_retries=5):
+    """Export model run results and retry if no results are returned."""
+    while (num_retries > 0):
+        task = model_run.export_v2(task_name, params=params)
+        assert task.name == task_name
+        task.wait_till_done()
+        assert task.status == "COMPLETE"
+        assert task.errors is None
+        task_results = task.result
+        if len(task_results) == 0:
+            num_retries -= 1
+            time.sleep(5)
+        else:
+            return task_results
+    return []
+
+
+def test_model_run_export_v2(model_run_with_data_rows):
+    model_run, labels = model_run_with_data_rows
+    label_ids = [label.uid for label in labels]
+    expected_data_rows = list(model_run.model_run_data_rows())
+
+    task_name = "test_task"
+    params = {"media_attributes": True, "predictions": True}
+    task_results = _model_run_export_v2_results(model_run, task_name, params)
+    assert len(task_results) == len(expected_data_rows)
+
+    for task_result in task_results:
+        # Check export param handling
+        assert 'media_attributes' in task_result and task_result[
+            'media_attributes'] is not None
+        exported_model_run = task_result['experiments'][
+            model_run.model_id]['runs'][model_run.uid]
+        task_label_ids_set = set(
+            map(lambda label: label['id'], exported_model_run['labels']))
+        task_prediction_ids_set = set(
+            map(lambda prediction: prediction['id'],
+                exported_model_run['predictions']))
+        for label_id in task_label_ids_set:
+            assert label_id in label_ids
+        for prediction_id in task_prediction_ids_set:
+            assert prediction_id in label_ids

tests/integration/export/legacy/test_export_model_run.py

@@ -0,0 +1,82 @@
+import json
+import time
+
+import pytest
+
+from labelbox import DataRow, ExportTask, StreamType
+
+
+class TestExportDataRow:
+
+    def test_with_data_row_object(self, client, data_row,
+                                  wait_for_data_row_processing):
+        data_row = wait_for_data_row_processing(client, data_row)
+        time.sleep(7)  # temp fix for ES indexing delay
+        export_task = DataRow.export(
+            client=client,
+            data_rows=[data_row],
+            task_name="TestExportDataRow:test_with_data_row_object",
+        )
+        export_task.wait_till_done()
+        assert export_task.status == "COMPLETE"
+        assert isinstance(export_task, ExportTask)
+        assert export_task.has_result()
+        assert export_task.has_errors() is False
+        assert export_task.get_total_file_size(
+            stream_type=StreamType.RESULT) > 0
+        assert export_task.get_total_lines(stream_type=StreamType.RESULT) == 1
+        assert (json.loads(list(export_task.get_stream())[0].json_str)
+                ["data_row"]["id"] == data_row.uid)
+
+    def test_with_id(self, client, data_row, wait_for_data_row_processing):
+        data_row = wait_for_data_row_processing(client, data_row)
+        time.sleep(7)  # temp fix for ES indexing delay
+        export_task = DataRow.export(client=client,
+                                     data_rows=[data_row.uid],
+                                     task_name="TestExportDataRow:test_with_id")
+        export_task.wait_till_done()
+        assert export_task.status == "COMPLETE"
+        assert isinstance(export_task, ExportTask)
+        assert export_task.has_result()
+        assert export_task.has_errors() is False
+        assert export_task.get_total_file_size(
+            stream_type=StreamType.RESULT) > 0
+        assert export_task.get_total_lines(stream_type=StreamType.RESULT) == 1
+        assert (json.loads(list(export_task.get_stream())[0].json_str)
+                ["data_row"]["id"] == data_row.uid)
+
+    def test_with_global_key(self, client, data_row,
+                             wait_for_data_row_processing):
+        data_row = wait_for_data_row_processing(client, data_row)
+        time.sleep(7)  # temp fix for ES indexing delay
+        export_task = DataRow.export(
+            client=client,
+            global_keys=[data_row.global_key],
+            task_name="TestExportDataRow:test_with_global_key",
+        )
+        export_task.wait_till_done()
+        assert export_task.status == "COMPLETE"
+        assert isinstance(export_task, ExportTask)
+        assert export_task.has_result()
+        assert export_task.has_errors() is False
+        assert export_task.get_total_file_size(
+            stream_type=StreamType.RESULT) > 0
+        assert export_task.get_total_lines(stream_type=StreamType.RESULT) == 1
+        assert (json.loads(list(export_task.get_stream())[0].json_str)
+                ["data_row"]["id"] == data_row.uid)
+
+    def test_with_invalid_id(self, client):
+        export_task = DataRow.export(
+            client=client,
+            data_rows=["invalid_id"],
+            task_name="TestExportDataRow:test_with_invalid_id",
+        )
+        export_task.wait_till_done()
+        assert export_task.status == "COMPLETE"
+        assert isinstance(export_task, ExportTask)
+        assert export_task.has_result() is False
+        assert export_task.has_errors() is False
+        assert export_task.get_total_file_size(
+            stream_type=StreamType.RESULT) is None
+        assert export_task.get_total_lines(
+            stream_type=StreamType.RESULT) is None

…gration/export_v2/test_export_project.py …ion/export/legacy/test_export_project.pytests/integration/export_v2/test_export_project.py renamed to tests/integration/export/legacy/test_export_project.py tests/integration/export_v2/test_export_project.py renamed to tests/integration/export/legacy/test_export_project.py

@@ -0,0 +1,78 @@
+import json
+
+import pytest
+
+from labelbox import ExportTask, StreamType
+
+
+class TestExportDataset:
+
+    @pytest.mark.parametrize("data_rows", [3], indirect=True)
+    def test_export(self, dataset, data_rows):
+        expected_data_row_ids = [dr.uid for dr in data_rows]
+
+        export_task = dataset.export(task_name="TestExportDataset:test_export")
+        export_task.wait_till_done()
+
+        assert export_task.status == "COMPLETE"
+        assert isinstance(export_task, ExportTask)
+        assert export_task.has_result()
+        assert export_task.has_errors() is False
+        assert export_task.get_total_file_size(
+            stream_type=StreamType.RESULT) > 0
+        assert export_task.get_total_lines(
+            stream_type=StreamType.RESULT) == len(expected_data_row_ids)
+        data_row_ids = list(
+            map(lambda x: json.loads(x.json_str)["data_row"]["id"],
+                export_task.get_stream()))
+        assert data_row_ids.sort() == expected_data_row_ids.sort()
+
+    @pytest.mark.parametrize("data_rows", [3], indirect=True)
+    def test_with_data_row_filter(self, dataset, data_rows):
+        datarow_filter_size = 3
+        expected_data_row_ids = [dr.uid for dr in data_rows
+                                ][:datarow_filter_size]
+        filters = {"data_row_ids": expected_data_row_ids}
+
+        export_task = dataset.export(
+            filters=filters,
+            task_name="TestExportDataset:test_with_data_row_filter")
+        export_task.wait_till_done()
+
+        assert export_task.status == "COMPLETE"
+        assert isinstance(export_task, ExportTask)
+        assert export_task.has_result()
+        assert export_task.has_errors() is False
+        assert export_task.get_total_file_size(
+            stream_type=StreamType.RESULT) > 0
+        assert export_task.get_total_lines(
+            stream_type=StreamType.RESULT) == datarow_filter_size
+        data_row_ids = list(
+            map(lambda x: json.loads(x.json_str)["data_row"]["id"],
+                export_task.get_stream()))
+        assert data_row_ids.sort() == expected_data_row_ids.sort()
+
+    @pytest.mark.parametrize("data_rows", [3], indirect=True)
+    def test_with_global_key_filter(self, dataset, data_rows):
+        datarow_filter_size = 2
+        expected_global_keys = [dr.global_key for dr in data_rows
+                               ][:datarow_filter_size]
+        filters = {"global_keys": expected_global_keys}
+
+        export_task = dataset.export(
+            filters=filters,
+            task_name="TestExportDataset:test_with_global_key_filter")
+        export_task.wait_till_done()
+
+        assert export_task.status == "COMPLETE"
+        assert isinstance(export_task, ExportTask)
+        assert export_task.has_result()
+        assert export_task.has_errors() is False
+        assert export_task.get_total_file_size(
+            stream_type=StreamType.RESULT) > 0
+        assert export_task.get_total_lines(
+            stream_type=StreamType.RESULT) == datarow_filter_size
+        global_keys = list(
+            map(lambda x: json.loads(x.json_str)["data_row"]["global_key"],
+                export_task.get_stream()))
+        assert global_keys.sort() == expected_global_keys.sort()