Fix tests

Author: Val Brodsky

labelbox/data/annotation_types/data/generic_data_row_data.py

@@ -11,8 +11,8 @@ class GenericDataRowData(BaseData, _NoCoercionMixin):
     url: Optional[str] = None
     class_name: Literal["GenericDataRowData"] = "GenericDataRowData"
 
-    def create_url(self, signer: Callable[[bytes], str]) -> None:
-        return None
+    def create_url(self, signer: Callable[[bytes], str]) -> Optional[str]:
+        return self.url
 
     @pydantic_compat.root_validator(pre=True)
     def validate_one_datarow_key_present(cls, data):

labelbox/data/annotation_types/label.py

@@ -53,15 +53,9 @@ class Label(pydantic_compat.BaseModel):
                             RelationshipAnnotation]] = []
     extra: Dict[str, Any] = {}
 
-    @staticmethod
-    def is_data_type(data: Union[Dict[str, Any], DataType]) -> bool:
-        if isinstance(data, DataType):
-            return True
-        return False
-
     @pydantic_compat.root_validator(pre=True)
     def validate_data(cls, label):
-        if not Label.is_data_type(label.get("data")):
+        if isinstance(label.get("data"), Dict):
             label["data"]["class_name"] = "GenericDataRowData"
         else:
             warnings.warn(

labelbox/schema/id_type.py

@@ -1,6 +1,6 @@
 import sys
 
-if sys.version_info > (3, 8):
+if sys.version_info >= (3, 9):
     from strenum import StrEnum
 
     class BaseStrEnum(StrEnum):

tests/conftest.py

@@ -121,13 +121,13 @@ def rest_url(environ: str) -> str:
     return 'http://host.docker.internal:8080/api/v1'
 
 
-def testing_api_key(environ: str) -> str:
-    for var in [
-            "LABELBOX_TEST_API_KEY_PROD", "LABELBOX_TEST_API_KEY_STAGING",
-            "LABELBOX_TEST_API_KEY_CUSTOM", "LABELBOX_TEST_API_KEY_LOCAL",
-            "LABELBOX_TEST_API_KEY"
-    ]:
-        value = os.environ.get(var)
+def testing_api_key(environ: Environ) -> str:
+    keys = [
+        f"LABELBOX_TEST_API_KEY_{environ.value.upper()}",
+        "LABELBOX_TEST_API_KEY"
+    ]
+    for key in keys:
+        value = os.environ.get(key)
         if value is not None:
             return value
     raise Exception("Cannot find API to use for tests")
@@ -147,7 +147,6 @@ def __init__(self, environ: str) -> None:
         api_url = graphql_url(environ)
         api_key = testing_api_key(environ)
         rest_endpoint = rest_url(environ)
-
         super().__init__(api_key,
                          api_url,
                          enable_experimental=True,

tests/data/annotation_import/conftest.py

@@ -1922,20 +1922,3 @@ def rename_cuid_key_recursive(d):
 @pytest.fixture
 def helpers():
     return Helpers
-
-
-@pytest.fixture
-def create_data_row_for_project(project, dataset, data_row_ndjson, batch_name):
-    data_row = dataset.create_data_row(data_row_ndjson)
-
-    project.create_batch(
-        batch_name,
-        [data_row.uid],  # sample of data row objects
-        5,  # priority between 1(Highest) - 5(lowest)
-    )
-    project.data_row_ids.append(data_row.uid)
-
-    yield data_row
-
-    data_row.delete()
-    project.delete()

tests/data/annotation_import/test_data_types.py

@@ -131,6 +131,19 @@ def get_annotation_comparison_dicts_from_export(export_result, data_row_id,
     return converted_annotations
 
 
+def create_data_row_for_project(project, dataset, data_row_ndjson, batch_name):
+    data_row = dataset.create_data_row(data_row_ndjson)
+
+    project.create_batch(
+        batch_name,
+        [data_row.uid],  # sample of data row objects
+        5,  # priority between 1(Highest) - 5(lowest)
+    )
+    project.data_row_ids.append(data_row.uid)
+
+    return data_row
+
+
 # TODO: Add VideoData. Currently label import job finishes without errors but project.export_labels() returns empty list.
 @pytest.mark.parametrize(
     "data_type_class",
@@ -147,10 +160,15 @@ def get_annotation_comparison_dicts_from_export(export_result, data_row_id,
         LlmResponseCreationData,
     ],
 )
-def test_import_data_types(client, configured_project, initial_dataset,
-                           rand_gen, data_row_json_by_data_type,
-                           annotations_by_data_type, data_type_class,
-                           create_data_row_for_project):
+def test_import_data_types(
+    client,
+    configured_project,
+    initial_dataset,
+    rand_gen,
+    data_row_json_by_data_type,
+    annotations_by_data_type,
+    data_type_class,
+):
     project = configured_project
     project_id = project.uid
     dataset = initial_dataset
@@ -193,7 +211,6 @@ def test_import_data_types_by_global_key(
     rand_gen,
     data_row_json_by_data_type,
     annotations_by_data_type,
-    create_data_row_for_project,
 ):
     project = configured_project
     project_id = project.uid
@@ -285,7 +302,6 @@ def test_import_data_types_v2(
     export_v2_test_helpers,
     rand_gen,
     helpers,
-    create_data_row_for_project,
 ):
     project = configured_project
     dataset = initial_dataset
@@ -355,7 +371,6 @@ def test_import_label_annotations(
     data_class,
     annotations,
     rand_gen,
-    create_data_row_for_project,
 ):
     project = configured_project_with_one_data_row
     dataset = initial_dataset

tests/data/annotation_import/test_generic_data_types.py

@@ -65,11 +65,28 @@
 ]
 
 
-def test_import_data_types_by_global_key(client, configured_project,
-                                         initial_dataset, rand_gen,
-                                         data_row_json_by_data_type,
-                                         annotations_by_data_type,
-                                         create_data_row_for_project):
+def create_data_row_for_project(project, dataset, data_row_ndjson, batch_name):
+    data_row = dataset.create_data_row(data_row_ndjson)
+
+    project.create_batch(
+        batch_name,
+        [data_row.uid],  # sample of data row objects
+        5,  # priority between 1(Highest) - 5(lowest)
+    )
+    project.data_row_ids.append(data_row.uid)
+
+    return data_row
+
+
+def test_import_data_types_by_global_key(
+    client,
+    configured_project,
+    initial_dataset,
+    rand_gen,
+    data_row_json_by_data_type,
+    annotations_by_data_type,
+    export_v2_test_helpers,
+):
     project = configured_project
     project_id = project.uid
     dataset = initial_dataset
@@ -93,17 +110,27 @@ def test_import_data_types_by_global_key(client, configured_project,
         ) for annotations in annotations_list
     ]
 
+    def find_data_row(dr):
+        return dr['data_row']['id'] == data_row.uid
+
     label_import = lb.LabelImport.create_from_objects(client, project_id,
                                                       f"test-import-image",
                                                       labels)
     label_import.wait_until_done()
 
     assert label_import.state == AnnotationImportState.FINISHED
     assert len(label_import.errors) == 0
-    exported_labels = project.export_labels(download=True)
-    objects = exported_labels[0]["Label"]["objects"]
-    classifications = exported_labels[0]["Label"]["classifications"]
+
+    result = export_v2_test_helpers.run_project_export_v2_task(project)
+    exported_data = list(filter(find_data_row, result))[0]
+    assert exported_data
+
+    label = exported_data['projects'][project.uid]['labels'][0]
+    annotations = label['annotations']
+    objects = annotations['objects']
+    classifications = annotations['classifications']
     assert len(objects) + len(classifications) == len(labels)
+
     data_row.delete()
 
 
@@ -151,11 +178,18 @@ def set_project_media_type_from_data_type(project, data_type_class):
         LlmResponseCreationData,
     ],
 )
-def test_import_data_types_v2(client, configured_project, initial_dataset,
-                              data_row_json_by_data_type,
-                              annotations_by_data_type_v2, data_type_class,
-                              exports_v2_by_data_type, export_v2_test_helpers,
-                              rand_gen, helpers, create_data_row_for_project):
+def test_import_data_types_v2(
+    client,
+    configured_project,
+    initial_dataset,
+    data_row_json_by_data_type,
+    annotations_by_data_type_v2,
+    data_type_class,
+    exports_v2_by_data_type,
+    export_v2_test_helpers,
+    rand_gen,
+    helpers,
+):
     project = configured_project
     dataset = initial_dataset
     project_id = project.uid