Update integration test

Author: Val Brodsky

labelbox/schema/id_type.py

@@ -0,0 +1,13 @@
+from enum import Enum
+
+
+class IdType(str, Enum):
+    """
+    The type of id used to identify a data row.
+    
+    Currently supported types are:
+        - DataRowId: The id assigned to a data row by Labelbox.
+        - GlobalKey: The id assigned to a data row by the user.
+    """
+    DataRowId = "ID"
+    GlobalKey = "GKEY"

labelbox/schema/project.py

@@ -30,7 +30,7 @@
 from labelbox.schema.export_params import ProjectExportParams
 from labelbox.schema.export_task import ExportTask
 from labelbox.schema.id_type import IdType
-from labelbox.schema.identifiable import DataRowIdentifier, UniqueId
+from labelbox.schema.identifiable import DataRowIdentifier, GlobalKey, UniqueId
 from labelbox.schema.identifiables import DataRowIdentifiers, UniqueIds
 from labelbox.schema.media_type import MediaType
 from labelbox.schema.queue_mode import QueueMode
@@ -62,16 +62,19 @@ def validate_labeling_parameter_overrides(
             )
         data_row_identifier = row[0]
         priority = row[1]
-        if not isinstance(data_row_identifier,
-                          Entity.DataRow) and not isinstance(
-                              data_row_identifier, DataRowIdentifier):
+        valid_types = (Entity.DataRow, UniqueId, GlobalKey)
+        if not isinstance(data_row_identifier, valid_types):
             raise TypeError(
-                f"Data row identifier should be be of type DataRow or Data Row Identifier. Found {type(data_row_identifier)}. Index: {idx}"
+                f"Data row identifier should be be of type DataRow, UniqueId or GlobalKey. Found {type(data_row_identifier)} for data_row_identifier {data_row_identifier}"
             )
 
         if not isinstance(priority, int):
+            if isinstance(data_row_identifier, Entity.DataRow):
+                id = data_row_identifier.uid
+            else:
+                id = data_row_identifier
             raise TypeError(
-                f"Priority must be an int. Found {type(priority)} for data_row_identifier {data_row_identifier}. Index: {idx}"
+                f"Priority must be an int. Found {type(priority)} for data_row_identifier {id}"
             )
 
 
@@ -1206,7 +1209,8 @@ def set_labeling_parameter_overrides(
         for data_row, priority in data:
             if isinstance(data_row, DataRow):
                 data_rows_with_identifiers += f"{{dataRowIdentifier: {{id: \"{data_row.uid}\", idType: {IdType.DataRowId}}}, priority: {priority}}},"
-            elif isinstance(data_row, DataRowIdentifier):
+            elif isinstance(data_row, UniqueId) or isinstance(
+                    data_row, GlobalKey):
                 data_rows_with_identifiers += f"{{dataRowIdentifier: {{id: \"{data_row.key}\", idType: {data_row.id_type}}}, priority: {priority}}},"
             else:
                 raise TypeError(

tests/integration/test_labeling_parameter_overrides.py

@@ -1,5 +1,6 @@
 import pytest
 from labelbox import DataRow
+from labelbox.schema.identifiable import GlobalKey, UniqueId
 from labelbox.schema.identifiables import GlobalKeys, UniqueIds
 
 
@@ -27,17 +28,36 @@ def test_labeling_parameter_overrides(consensus_project_with_batch):
     for override in updated_overrides:
         assert isinstance(override.data_row(), DataRow)
 
+    data = [(UniqueId(data_rows[0].uid), 1, 2), (UniqueId(data_rows[1].uid), 2),
+            (UniqueId(data_rows[2].uid), 3)]
+    success = project.set_labeling_parameter_overrides(data)
+    assert success
+    updated_overrides = list(project.labeling_parameter_overrides())
+    assert len(updated_overrides) == 3
+    assert {o.number_of_labels for o in updated_overrides} == {1, 1, 1}
+    assert {o.priority for o in updated_overrides} == {1, 2, 3}
+
+    data = [(GlobalKey(data_rows[0].global_key), 2, 2),
+            (GlobalKey(data_rows[1].global_key), 3, 3),
+            (GlobalKey(data_rows[2].global_key), 4)]
+    success = project.set_labeling_parameter_overrides(data)
+    assert success
+    updated_overrides = list(project.labeling_parameter_overrides())
+    assert len(updated_overrides) == 3
+    assert {o.number_of_labels for o in updated_overrides} == {1, 1, 1}
+    assert {o.priority for o in updated_overrides} == {2, 3, 4}
+
     with pytest.raises(TypeError) as exc_info:
         data = [(data_rows[2], "a_string", 3)]
         project.set_labeling_parameter_overrides(data)
     assert str(exc_info.value) == \
-           f"Priority must be an int. Found <class 'str'> for data_row {data_rows[2]}. Index: 0"
+            f"Priority must be an int. Found <class 'str'> for data_row_identifier {data_rows[2].uid}"
 
     with pytest.raises(TypeError) as exc_info:
         data = [(data_rows[2].uid, 1)]
         project.set_labeling_parameter_overrides(data)
     assert str(exc_info.value) == \
-           "data_row should be be of type DataRow. Found <class 'str'>. Index: 0"
+            f"Data row identifier should be be of type DataRow, UniqueId or GlobalKey. Found <class 'str'> for data_row_identifier {data_rows[2].uid}"
 
 
 def test_set_labeling_priority(consensus_project_with_batch):