Update tests

Author: mnoszczak

labelbox/schema/export_params.py

@@ -1,5 +1,3 @@
-import sys
-
 from typing import Optional
 if sys.version_info >= (3, 8):
     from typing import TypedDict

labelbox/schema/project.py

@@ -20,7 +20,7 @@
 from labelbox.pagination import PaginatedCollection
 from labelbox.schema.consensus_settings import ConsensusSettings
 from labelbox.schema.data_row import DataRow
-from labelbox.schema.export_params import ModelRunExportParams
+from labelbox.schema.export_params import ProjectExportParams
 from labelbox.schema.media_type import MediaType
 from labelbox.schema.queue_mode import QueueMode
 from labelbox.schema.resource_tag import ResourceTag
@@ -287,7 +287,7 @@ def label_generator(self, timeout_seconds=600, **kwargs):
             return LBV1Converter.deserialize_video(json_data, self.client)
         return LBV1Converter.deserialize(json_data)
 
-    def export_labels_v2(
+    def export_labels(
             self,
             download=False,
             timeout_seconds=1800,
@@ -379,23 +379,22 @@ def _validate_datetime(string_date: str) -> bool:
     """
     Creates a project run export task with the given params and returns the task.
     
-    >>>    export_task = export_labels_v2("my_export_task", filter={"media_attributes": True})
+    >>>    export_task = export_v2("my_export_task", filter={"media_attributes": True})
     
     """
 
-    def export_labels_v2(self, task_name: str,
-                         params: Optional[ModelRunExportParams]) -> Task:
+    def export_v2(self, task_name: str,
+                         params: Optional[ProjectExportParams]) -> Task:
         _params = params or {}
-        mutation_name = "exportDataRows"
-        create_task_query_str = """mutation exportDataRowsPyApi($input: ExportDataRowsInput!){
+        mutation_name = "exportDataRowsInProject"
+        create_task_query_str = """mutation exportDataRowsInProjectPyApi($input: ExportDataRowsInProjectInput!){
           %s(input: $input) {taskId} }
           """ % (mutation_name)
         params = {
             "input": {
                 "taskName": task_name,
                 "filters": {
-                    "modelRunIds": [],
-                    "projectIds": [self.uid]
+                    "projectId": self.uid
                 },
                 "params": {
                     "includeAttachments":
@@ -412,9 +411,6 @@ def export_labels_v2(self, task_name: str,
                         _params.get('include_labels', False),
                     "includePerformanceDetails":
                         _params.get('include_performance_details', False),
-                    # Arguments locked based on exectuion context
-                    "includeModelRuns":
-                        False,
                 },
             }
         }

tests/integration/test_label.py

@@ -1,3 +1,4 @@
+import json
 import time
 
 import pytest
@@ -43,18 +44,47 @@ def test_label_export(configured_project_with_label):
     # The new exporter doesn't work with the create_label mutation
 
 
-def test_label_export_v2(configured_project_with_label):
+def test_export_v2(configured_project_with_label):
     project, _, _, label = configured_project_with_label
     label_id = label.uid
     # Wait for exporter to retrieve latest labels
     time.sleep(10)
     task_name = "test_label_export_v2"
 
-    task = project.export_labels_v2(task_name, filter={"project_details": True})
+    # TODO: Right now we don't have a way to test this
+    include_performance_details = True
+    task = project.export_v2(
+        task_name,
+        params={
+            "include_performance_details": include_performance_details,
+            "include_labels": True
+        })
     assert task.name == task_name
     task.wait_till_done()
     assert task.status == "COMPLETE"
-    # TODO: Download result and check it
+
+    def download_result(result_url):
+        response = requests.get(result_url)
+        response.raise_for_status()
+        data = [json.loads(line) for line in response.text.splitlines()]
+        return data
+
+    task_results = download_result(task.result_url)
+
+    for task_result in task_results:
+        assert len(task_result['errors']) == 0
+        task_project = task_result['projects'][project.uid]
+        task_project_label_ids_set = set(
+            map(lambda prediction: prediction['id'], task_project['labels']))
+        assert label_id in task_project_label_ids_set
+
+        # TODO: Add back in when we have a way to test this
+        # if include_performance_details:
+        #     assert 'include_performance_details' in task_result and task_result[
+        #         'include_performance_details'] is not None
+        # else:
+        #     assert 'include_performance_details' not in task_result or task_result[
+        #         'include_performance_details'] is None
 
 
 # TODO: Skipping this test in staging due to label not updating