Fixing tests

Author: Val Brodsky

libs/labelbox/src/labelbox/schema/dataset.py

@@ -283,9 +283,14 @@ def create_data_rows(
 
         NOTE  dicts and strings items can not be mixed in the same call. It is a responsibility of the caller to ensure that all items are of the same type.
         """
-        if isinstance(items[0], str):
-            items = self._build_from_local_paths(items)  # Assume list of file paths
-        specs = DataRowCreateItem.build(self.uid, items)
+        string_items = [item for item in items if isinstance(item, str)]
+        dict_items = [item for item in items if isinstance(item, dict)]
+        dict_string_items = []
+
+        if len(string_items) > 0:
+            dict_string_items = self._build_from_local_paths(string_items)
+        specs = DataRowCreateItem.build(self.uid,
+                                        dict_items + dict_string_items)
         return self._exec_upsert_data_rows(specs, file_upload_thread_count)
 
     def _build_from_local_paths(

libs/labelbox/src/labelbox/schema/internal/data_row_create_upsert.py

@@ -25,13 +25,12 @@ def build(
             key = item.pop('key', None)
             if not key:
                 key = {'type': 'AUTO', 'value': ''}
-            elif isinstance(key, key_types):
+            elif isinstance(key, key_types):  # type: ignore
                 key = {'type': key.id_type.value, 'value': key.key}
             else:
                 if not key_types:
                     raise ValueError(
-                        f"Can not have a key for this item, got: {key}"
-                    )
+                        f"Can not have a key for this item, got: {key}")
                 raise ValueError(
                     f"Key must be an instance of {', '.join([t.__name__ for t in key_types])}, got: {type(item['key']).__name__}"
                 )
@@ -53,14 +52,22 @@ def is_empty(self) -> bool:
 class DataRowUpsertItem(DataRowItemBase):
 
     @classmethod
-    def build(cls, dataset_id: str,
-              items: List[dict]) -> List["DataRowUpsertItem"]:
+    def build(
+        cls,
+        dataset_id: str,
+        items: List[dict],
+        key_types: Optional[Tuple[type, ...]] = ()
+    ) -> List["DataRowItemBase"]:
         return super().build(dataset_id, items, (UniqueId, GlobalKey))
 
 
 class DataRowCreateItem(DataRowItemBase):
 
     @classmethod
-    def build(cls, dataset_id: str,
-              items: List[dict]) -> List["DataRowCreateItem"]:
+    def build(
+        cls,
+        dataset_id: str,
+        items: List[dict],
+        key_types: Optional[Tuple[type, ...]] = ()
+    ) -> List["DataRowItemBase"]:
         return super().build(dataset_id, items, ())

libs/labelbox/src/labelbox/schema/internal/datarow_upload_constants.py

@@ -1,3 +1,4 @@
 MAX_DATAROW_PER_API_OPERATION = 150_000
 FILE_UPLOAD_THREAD_COUNT = 20
 UPSERT_CHUNK_SIZE = 10_000
+DOWNLOAD_RESULT_PAGE_SIZE = 5_000

libs/labelbox/src/labelbox/schema/task.py

@@ -10,7 +10,10 @@
 from labelbox.orm.model import Field, Relationship, Entity
 
 from labelbox.pagination import PaginatedCollection
-from labelbox.schema.internal.datarow_upload_constants import MAX_DATAROW_PER_API_OPERATION
+from labelbox.schema.internal.datarow_upload_constants import (
+    MAX_DATAROW_PER_API_OPERATION,
+    DOWNLOAD_RESULT_PAGE_SIZE,
+)
 
 if TYPE_CHECKING:
     from labelbox import User
@@ -52,6 +55,10 @@ class Task(DbObject):
     created_by = Relationship.ToOne("User", False, "created_by")
     organization = Relationship.ToOne("Organization")
 
+    def __eq__(self, task):
+        return isinstance(
+            task, Task) and task.uid == self.uid and task.type == self.type
+
     # Import and upsert have several instances of special casing
     def is_creation_task(self) -> bool:
         return self.name == 'JSON Import' or self.type == 'adv-upsert-data-rows'
@@ -227,21 +234,23 @@ def __init__(self, *args, **kwargs):
         self._user = None
 
     @property
-    def result(self) -> Union[List[Dict[str, Any]]]:
+    def result(self) -> Optional[List[Dict[str, Any]]]:  # type: ignore
         if self.status == "FAILED":
             raise ValueError(f"Job failed. Errors : {self.errors}")
         return self._results_as_list()
 
     @property
-    def errors(self) -> Optional[Dict[str, Any]]:
+    def errors(self) -> Optional[List[Dict[str, Any]]]:  # type: ignore
         return self._errors_as_list()
 
     @property
-    def created_data_rows(self) -> Optional[Dict[str, Any]]:
+    def created_data_rows(  # type: ignore
+            self) -> Optional[List[Dict[str, Any]]]:
         return self.result
 
     @property
-    def failed_data_rows(self) -> Optional[Dict[str, Any]]:
+    def failed_data_rows(  # type: ignore
+            self) -> Optional[List[Dict[str, Any]]]:
         return self.errors
 
     @property
@@ -253,7 +262,7 @@ def errors_all(self) -> PaginatedCollection:
         return self._download_errors_paginated()
 
     def _download_results_paginated(self) -> PaginatedCollection:
-        page_size = 900  # hardcode to avoid overloading the server
+        page_size = DOWNLOAD_RESULT_PAGE_SIZE
         from_cursor = None
 
         query_str = """query SuccessesfulDataRowImportsPyApi($taskId: ID!, $first: Int, $from: String)  {
@@ -292,7 +301,7 @@ def _download_results_paginated(self) -> PaginatedCollection:
         )
 
     def _download_errors_paginated(self) -> PaginatedCollection:
-        page_size = 5000  # hardcode to avoid overloading the server
+        page_size = DOWNLOAD_RESULT_PAGE_SIZE  # hardcode to avoid overloading the server
         from_cursor = None
 
         query_str = """query FailedDataRowImportsPyApi($taskId: ID!, $first: Int, $from: String)  {
@@ -306,6 +315,16 @@ def _download_errors_paginated(self) -> PaginatedCollection:
                                     externalId
                                     globalKey
                                     rowData
+                                        metadata {
+                                            schemaId
+                                            value
+                                            name
+                                        }
+                                        attachments {
+                                            type
+                                            value
+                                            name
+                                        }                                    
                                 }
                             }
                         }
@@ -318,28 +337,30 @@ def _download_errors_paginated(self) -> PaginatedCollection:
             'from': from_cursor,
         }
 
+        def convert_errors_to_legacy_format(client, data_row):
+            spec = data_row.get('spec', {})
+            return {
+                'message':
+                    data_row.get('message'),
+                'failedDataRows': [{
+                    'externalId': spec.get('externalId'),
+                    'rowData': spec.get('rowData'),
+                    'globalKey': spec.get('globalKey'),
+                    'metadata': spec.get('metadata', []),
+                    'attachments': spec.get('attachments', []),
+                }]
+            }
+
         return PaginatedCollection(
             client=self.client,
             query=query_str,
             params=params,
             dereferencing=['failedDataRowImports', 'results'],
-            obj_class=lambda _, data_row: {
-                'error':
-                    data_row.get('message'),
-                'external_id':
-                    data_row.get('spec').get('externalId')
-                    if data_row.get('spec') else None,
-                'row_data':
-                    data_row.get('spec').get('rowData')
-                    if data_row.get('spec') else None,
-                'global_key':
-                    data_row.get('spec').get('globalKey')
-                    if data_row.get('spec') else None,
-            },
+            obj_class=convert_errors_to_legacy_format,
             cursor_path=['failedDataRowImports', 'after'],
         )
 
-    def _results_as_list(self) -> List[Dict[str, Any]]:
+    def _results_as_list(self) -> Optional[List[Dict[str, Any]]]:
         total_downloaded = 0
         results = []
         data = self._download_results_paginated()
@@ -350,9 +371,12 @@ def _results_as_list(self) -> List[Dict[str, Any]]:
             if total_downloaded >= self.__max_donwload_size:
                 break
 
+        if len(results) == 0:
+            return None
+
         return results
 
-    def _errors_as_list(self) -> List[Dict[str, Any]]:
+    def _errors_as_list(self) -> Optional[List[Dict[str, Any]]]:
         total_downloaded = 0
         errors = []
         data = self._download_errors_paginated()
@@ -363,4 +387,7 @@ def _errors_as_list(self) -> List[Dict[str, Any]]:
             if total_downloaded >= self.__max_donwload_size:
                 break
 
+        if len(errors) == 0:
+            return None
+
         return errors

libs/labelbox/tests/integration/test_data_rows.py

@@ -238,7 +238,7 @@ def test_data_row_bulk_creation_from_file(dataset, local_image_file, image_url):
         assert task.has_errors() is False
         results = [r for r in task.result_all]
         row_data = [result["row_data"] for result in results]
-        assert row_data == [image_url, image_url]
+        assert len(row_data) == 2
 
 
 def test_data_row_bulk_creation_from_row_data_file_external_id(
@@ -252,12 +252,14 @@ def test_data_row_bulk_creation_from_row_data_file_external_id(
             "row_data": image_url,
             'external_id': 'some_name2'
         }])
+        task.wait_till_done()
         assert task.status == "COMPLETE"
         assert len(task.result) == 2
         assert task.has_errors() is False
         results = [r for r in task.result_all]
         row_data = [result["row_data"] for result in results]
-        assert row_data == [image_url, image_url]
+        assert len(row_data) == 2
+        assert image_url in row_data
 
 
 def test_data_row_bulk_creation_from_row_data_file(dataset, rand_gen,
@@ -275,7 +277,7 @@ def test_data_row_bulk_creation_from_row_data_file(dataset, rand_gen,
         assert task.has_errors() is False
         results = [r for r in task.result_all]
         row_data = [result["row_data"] for result in results]
-        assert row_data == [image_url, image_url]
+        assert len(row_data) == 2
 
 
 @pytest.mark.slow
@@ -899,6 +901,7 @@ def test_create_data_rows_result(client, dataset, image_url):
             DataRow.external_id: "row1",
         },
     ])
+    task.wait_till_done()
     assert task.errors is None
     for result in task.result:
         client.get_data_row(result['id'])
@@ -973,8 +976,16 @@ def test_data_row_bulk_creation_with_same_global_keys(dataset, sample_image,
         'message'] == f"Duplicate global key: '{global_key_1}'"
     assert task.failed_data_rows[0]['failedDataRows'][0][
         'externalId'] == sample_image
-    assert task.created_data_rows[0]['externalId'] == sample_image
-    assert task.created_data_rows[0]['globalKey'] == global_key_1
+    assert task.created_data_rows[0]['external_id'] == sample_image
+    assert task.created_data_rows[0]['global_key'] == global_key_1
+
+    errors = task.errors_all
+    all_errors = [er for er in errors]
+    assert len(all_errors) == 1
+    assert task.has_errors() is True
+
+    all_results = [result for result in task.result_all]
+    assert len(all_results) == 1
 
 
 def test_data_row_delete_and_create_with_same_global_key(

libs/labelbox/tests/integration/test_data_rows_upsert.py

@@ -208,9 +208,8 @@ def test_multiple_chunks(self, client, dataset, image_url):
         mocked_chunk_size = 3
         with patch('labelbox.client.Client.upload_data',
                    wraps=client.upload_data) as spy_some_function:
-            with patch(
-                    'labelbox.schema.dataset.Dataset._Dataset__upsert_chunk_size',
-                    new=mocked_chunk_size):
+            with patch('labelbox.schema.dataset.UPSERT_CHUNK_SIZE',
+                       new=mocked_chunk_size):
                 task = dataset.upsert_data_rows([{
                     'row_data': image_url
                 } for i in range(10)])

libs/labelbox/tests/integration/test_dataset.py

@@ -3,7 +3,7 @@
 from labelbox import Dataset
 from labelbox.exceptions import ResourceNotFoundError, MalformedQueryException, InvalidQueryError
 from labelbox.schema.dataset import MAX_DATAROW_PER_API_OPERATION
-from labelbox.schema.internal.datarow_uploader import DataRowUploader
+from labelbox.schema.internal.data_row_uploader import DataRowUploader
 
 
 def test_dataset(client, rand_gen):
@@ -166,12 +166,3 @@ def test_create_descriptor_file(dataset):
             'content_type': 'application/json',
             'filename': 'json_import.json'
         }
-
-
-def test_max_dataset_datarow_upload(dataset, image_url, rand_gen):
-    external_id = str(rand_gen)
-    items = [dict(row_data=image_url, external_id=external_id)
-            ] * (MAX_DATAROW_PER_API_OPERATION + 1)
-
-    with pytest.raises(MalformedQueryException):
-        dataset.create_data_rows(items)

libs/labelbox/tests/integration/test_task.py

@@ -61,11 +61,12 @@ def test_task_success_json(dataset, image_url, snapshot):
 @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
-    project.export_labels()
+    # TODO: Move to export_v2
+    res = project.export_labels()
     user = client.get_user()
     task = None
     for task in user.created_tasks():
-        if task.name != 'JSON Import':
+        if task.name != 'JSON Import' and task.type != 'adv-upsert-data-rows':
             break
 
     with pytest.raises(ValueError) as exc_info: