Update tests

Author: mnoszczak

labelbox/schema/export_params.py

@@ -1,25 +1,23 @@
 import sys
 
 from typing import Optional
-if sys.version_info >= (3, 8):
-    from typing import TypedDict
-else:
-    from typing_extensions import TypedDict
 
+from pydantic import BaseModel
 
-class DataRowParams(TypedDict):
-    include_data_row_details: Optional[bool]
-    include_media_attributes: Optional[bool]
-    include_metadata_fields: Optional[bool]
-    include_attachments: Optional[bool]
 
+class DataRowParams(BaseModel):
+    include_data_row_details: Optional[bool] = None
+    include_media_attributes: Optional[bool] = None
+    include_metadata_fields: Optional[bool] = None
+    include_attachments: Optional[bool] = None
 
-class ProjectExportParams(DataRowParams):
-    include_project_details: Optional[bool]
-    include_label_details: Optional[bool]
-    include_performance_details: Optional[bool]
 
+class ProjectExportParams(BaseModel):
+    include_project_details: Optional[bool] = None
+    include_label_details: Optional[bool] = None
+    include_performance_details: Optional[bool] = None
 
-class ModelRunExportParams(DataRowParams):
+
+class ModelRunExportParams(BaseModel):
     # TODO: Add model run fields
     pass

labelbox/schema/model_run.py

@@ -14,7 +14,7 @@
 from labelbox.orm.db_object import DbObject, experimental
 from labelbox.schema.export_params import ModelRunExportParams
 from labelbox.schema.task import Task
-from labelbox.schema.user import User  # type: ignore
+from labelbox.schema.user import User
 
 if TYPE_CHECKING:
     from labelbox import MEAPredictionImport
@@ -450,25 +450,24 @@ def export_labels(
             time.sleep(sleep_time)
 
     """
-    Creates a model run export task with the given filter and returns the task.
+    Creates a model 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", params={"media_attributes": True})
     
     """
 
-    def export_labels_v2(self, task_name: str,
-                         params: Optional[ModelRunExportParams]) -> Task:
+    def export_v2(self, task_name: str,
+                  params: Optional[ModelRunExportParams]) -> Task:
         _params = params or {}
-        mutation_name = "exportDataRows"
-        create_task_query_str = """mutation exportDataRowsPyApi($input: ExportDataRowsInput!){
+        mutation_name = "exportDataRowsInModelRun"
+        create_task_query_str = """mutation exportDataRowsInModelRunPyApi($input: ExportDataRowsInModelRunInput!){
           %s(input: $input) {taskId} }
           """ % (mutation_name)
         params = {
             "input": {
                 "taskName": task_name,
                 "filters": {
-                    "modelRunIds": [self.uid],
-                    "projectIds": []
+                    "modelRunId": self.uid
                 },
                 "params": {
                     "includeAttachments":
@@ -480,8 +479,6 @@ def export_labels_v2(self, task_name: str,
                     "includeDataRowDetails":
                         _params.get('include_data_row_details', False),
                     # Arguments locked based on exectuion context
-                    "includeModelRuns":
-                        True,
                     "includeProjectDetails":
                         False,
                     "includeLabels":

tests/integration/annotation_import/test_model_run.py

@@ -1,8 +1,11 @@
+import json
 import time
 import os
 import pytest
 
 from collections import Counter
+
+import requests
 from labelbox import DataSplit, ModelRun
 
 
@@ -114,15 +117,50 @@ def test_model_run_export_labels(model_run_with_model_run_data_rows):
     assert len(labels) == 3
 
 
-def test_model_run_export_labels_v2(model_run_with_model_run_data_rows):
+def test_model_run_export_v2(model_run_with_model_run_data_rows,
+                             configured_project):
     task_name = "test_task"
-    params = {"media_attributes": True}
-    task = model_run_with_model_run_data_rows.export_labels_v2(task_name,
-                                                               params=params)
+
+    media_attributes = True
+    params = {"media_attributes": media_attributes}
+    task = model_run_with_model_run_data_rows.export_v2(task_name,
+                                                        params=params)
     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)
+
+    label_ids = [label.uid for label in configured_project.labels()]
+    label_ids_set = set(label_ids)
+
+    assert len(task_results) == len(label_ids)
+    for task_result in task_results:
+        assert len(task_result['errors']) == 0
+        # Check export param handling
+        if media_attributes:
+            assert 'media_attributes' in task_result and task_result[
+                'media_attributes'] is not None
+        else:
+            assert 'media_attributes' not in task_result or task_result[
+                'media_attributes'] is None
+        model_run = task_result['models'][
+            model_run_with_model_run_data_rows.model_id]['model_runs'][
+                model_run_with_model_run_data_rows.uid]
+        task_label_ids_set = set(
+            map(lambda label: label['id'], model_run['labels']))
+        task_prediction_ids_set = set(
+            map(lambda prediction: prediction['id'], model_run['predictions']))
+        for label_id in task_label_ids_set:
+            assert label_id in label_ids_set
+        for prediction_id in task_prediction_ids_set:
+            assert prediction_id in label_ids_set
 
 
 @pytest.mark.skipif(condition=os.environ['LABELBOX_TEST_ENVIRON'] == "onprem",