Adding tests

Author: Kevin Kim

labelbox/schema/dataset.py

@@ -82,14 +82,8 @@ def create_data_row(self, **kwargs) -> "DataRow":
 
         # Parse metadata fields, if they are provided
         if DataRow.custom_metadata.name in kwargs:
-            mdo = self.client.get_data_row_metadata_ontology()
-            metadata_fields = kwargs[DataRow.custom_metadata.name]
-            metadata = list(
-                chain.from_iterable(
-                    mdo.parse_upsert(m) for m in metadata_fields))
-            kwargs[DataRow.custom_metadata.name] = [
-                md.dict(by_alias=True) for md in metadata
-            ]
+            kwargs[DataRow.custom_metadata.name] = self._parse_metadata(
+                kwargs[DataRow.custom_metadata.name])
 
         return self.client._create(DataRow, kwargs)
 
@@ -269,32 +263,10 @@ def validate_attachments(item):
                     )
             return attachments
 
-        def convert_metadata_field(metadata_field):
-            if isinstance(metadata_field, DataRowMetadataField):
-                return metadata_field
-            elif isinstance(metadata_field, dict):
-                if not all (key in metadata_field for key in ("schema_id", "value")):
-                    raise ValueError(f"Custom metadata fields must have 'schema_id' and 'value' keys")
-                return DataRowMetadataField(schema_id=metadata_field["schema_id"], value=metadata_field["value"])
-            else:
-                raise ValueError(f"Metadata field is neither 'DataRowMetadataField' type or a dictionary!")
-
         def parse_metadata_fields(item):
             metadata_fields = item.get('custom_metadata')
             if metadata_fields:
-                # Convert all metadata fields to DataRowMetadataField type
-                metadata_fields = [
-                    convert_metadata_field(m)
-                    for m in metadata_fields
-                ]
-                mdo = self.client.get_data_row_metadata_ontology()
-                metadata = list(
-                    chain.from_iterable(
-                        mdo.parse_upsert(m) for m in metadata_fields))
-                metadata_fields = [
-                    md.dict(by_alias=True) for md in metadata
-                ]
-                item['custom_metadata'] = metadata_fields
+                item['custom_metadata'] = self._parse_metadata(metadata_fields)
 
         def format_row(item):
             # Formats user input into a consistent dict structure
@@ -329,7 +301,7 @@ def convert_item(item):
             # Don't make any changes to tms data
             if "tileLayerUrl" in item:
                 validate_attachments(item)
-                return item                
+                return item
             # Convert all payload variations into the same dict format
             item = format_row(item)
             # Make sure required keys exist (and there are no extra keys)
@@ -455,3 +427,28 @@ def export_data_rows(self, timeout_seconds=120) -> Generator:
             logger.debug("Dataset '%s' data row export, waiting for server...",
                          self.uid)
             time.sleep(sleep_time)
+
+    def _convert_metadata_field(self, metadata_field):
+        if isinstance(metadata_field, DataRowMetadataField):
+            return metadata_field
+        elif isinstance(metadata_field, dict):
+            if not all(key in metadata_field for key in ("schema_id", "value")):
+                raise ValueError(
+                    f"Custom metadata field '{metadata_field}' must have 'schema_id' and 'value' keys"
+                )
+            return DataRowMetadataField(schema_id=metadata_field["schema_id"],
+                                        value=metadata_field["value"])
+        else:
+            raise ValueError(
+                f"Metadata field '{metadata_field}' is neither 'DataRowMetadataField' type or a dictionary!"
+            )
+
+    def _parse_metadata(self, metadata_fields):
+        # Convert all metadata fields to DataRowMetadataField type
+        metadata_fields = [
+            self._convert_metadata_field(m) for m in metadata_fields
+        ]
+        mdo = self.client.get_data_row_metadata_ontology()
+        parsed_metadata = list(
+            chain.from_iterable(mdo.parse_upsert(m) for m in metadata_fields))
+        return [m.dict(by_alias=True) for m in parsed_metadata]

tests/integration/test_data_rows.py

@@ -15,6 +15,10 @@
 EMBEDDING_SCHEMA_ID = "ckpyije740000yxdk81pbgjdc"
 TEXT_SCHEMA_ID = "cko8s9r5v0001h2dk9elqdidh"
 CAPTURE_DT_SCHEMA_ID = "cko8sdzv70006h2dk8jg64zvb"
+EXPECTED_METADATA_SCHEMA_IDS = [
+    SPLIT_SCHEMA_ID, TEST_SPLIT_ID, EMBEDDING_SCHEMA_ID, TEXT_SCHEMA_ID,
+    CAPTURE_DT_SCHEMA_ID
+].sort()
 
 
 def make_metadata_fields():
@@ -31,6 +35,27 @@ def make_metadata_fields():
     return fields
 
 
+def make_metadata_fields_dict():
+    embeddings = [0.0] * 128
+    msg = "A message"
+    time = datetime.utcnow()
+
+    fields = [{
+        "schema_id": SPLIT_SCHEMA_ID,
+        "value": TEST_SPLIT_ID
+    }, {
+        "schema_id": CAPTURE_DT_SCHEMA_ID,
+        "value": time
+    }, {
+        "schema_id": TEXT_SCHEMA_ID,
+        "value": msg
+    }, {
+        "schema_id": EMBEDDING_SCHEMA_ID,
+        "value": embeddings
+    }]
+    return fields
+
+
 def test_get_data_row(datarow, client):
     assert client.get_data_row(datarow.uid)
 
@@ -152,7 +177,7 @@ def test_data_row_single_creation(dataset, rand_gen, image_url):
         assert requests.get(data_row_2.row_data).content == data
 
 
-def test_data_row_single_creation_with_metadata(dataset, rand_gen, image_url):
+def test_create_data_row_with_metadata(dataset, image_url):
     client = dataset.client
     assert len(list(dataset.data_rows())) == 0
 
@@ -167,38 +192,133 @@ def test_data_row_single_creation_with_metadata(dataset, rand_gen, image_url):
         requests.get(data_row.row_data).content
     assert data_row.media_attributes is not None
     assert len(data_row.custom_metadata) == 5
+    assert [m["schemaId"] for m in data_row.custom_metadata
+           ].sort() == EXPECTED_METADATA_SCHEMA_IDS
 
-    with NamedTemporaryFile() as fp:
-        data = rand_gen(str).encode()
-        fp.write(data)
-        fp.flush()
-        data_row_2 = dataset.create_data_row(row_data=fp.name)
-        assert len(list(dataset.data_rows())) == 2
-        assert requests.get(data_row_2.row_data).content == data
 
+def test_create_data_row_with_metadata_dict(dataset, image_url):
+    client = dataset.client
+    assert len(list(dataset.data_rows())) == 0
 
-def test_data_row_single_creation_with_invalid_metadata(dataset, image_url):
+    data_row = dataset.create_data_row(
+        row_data=image_url, custom_metadata=make_metadata_fields_dict())
 
-    def make_invalid_metadata_fields():
-        embeddings = [0.0] * 128
-        msg = "A message"
-        time = datetime.utcnow()
+    assert len(list(dataset.data_rows())) == 1
+    assert data_row.dataset() == dataset
+    assert data_row.created_by() == client.get_user()
+    assert data_row.organization() == client.get_organization()
+    assert requests.get(image_url).content == \
+        requests.get(data_row.row_data).content
+    assert data_row.media_attributes is not None
+    assert len(data_row.custom_metadata) == 5
+    assert [m["schemaId"] for m in data_row.custom_metadata
+           ].sort() == EXPECTED_METADATA_SCHEMA_IDS
 
-        fields = [
-            DataRowMetadataField(schema_id=SPLIT_SCHEMA_ID,
-                                 value=TEST_SPLIT_ID),
-            DataRowMetadataField(schema_id=CAPTURE_DT_SCHEMA_ID, value=time),
-            DataRowMetadataField(schema_id=TEXT_SCHEMA_ID, value=msg),
-            DataRowMetadataField(schema_id=EMBEDDING_SCHEMA_ID,
-                                 value=embeddings),
-            DataRowMetadataField(schema_id=EMBEDDING_SCHEMA_ID,
-                                 value=embeddings),
-        ]
-        return fields
+
+def test_create_data_row_with_invalid_metadata(dataset, image_url):
+    fields = make_metadata_fields()
+    fields.append(
+        DataRowMetadataField(schema_id=EMBEDDING_SCHEMA_ID, value=[0.0] * 128))
 
     with pytest.raises(labelbox.exceptions.MalformedQueryException) as excinfo:
-        dataset.create_data_row(row_data=image_url,
-                                custom_metadata=make_invalid_metadata_fields())
+        dataset.create_data_row(row_data=image_url, custom_metadata=fields)
+
+
+def test_create_data_rows_with_metadata(dataset, image_url):
+    client = dataset.client
+    assert len(list(dataset.data_rows())) == 0
+
+    task = dataset.create_data_rows([
+        {
+            DataRow.row_data: image_url,
+            DataRow.external_id: "row1",
+            DataRow.custom_metadata: make_metadata_fields()
+        },
+        {
+            DataRow.row_data: image_url,
+            DataRow.external_id: "row2",
+            "custom_metadata": make_metadata_fields()
+        },
+        {
+            DataRow.row_data: image_url,
+            DataRow.external_id: "row3",
+            DataRow.custom_metadata: make_metadata_fields_dict()
+        },
+        {
+            DataRow.row_data: image_url,
+            DataRow.external_id: "row4",
+            "custom_metadata": make_metadata_fields_dict()
+        },
+    ])
+    task.wait_till_done()
+
+    assert len(list(dataset.data_rows())) == 4
+    for r in ["row1", "row2", "row3", "row4"]:
+        row = list(dataset.data_rows(where=DataRow.external_id == r))[0]
+        assert row.dataset() == dataset
+        assert row.created_by() == client.get_user()
+        assert row.organization() == client.get_organization()
+        assert requests.get(image_url).content == \
+            requests.get(row.row_data).content
+        assert row.media_attributes is not None
+        assert len(row.custom_metadata) == 5
+        assert [m["schemaId"] for m in row.custom_metadata
+               ].sort() == EXPECTED_METADATA_SCHEMA_IDS
+
+
+def test_create_data_rows_with_invalid_metadata(dataset, image_url):
+    fields = make_metadata_fields()
+    fields.append(
+        DataRowMetadataField(schema_id=EMBEDDING_SCHEMA_ID, value=[0.0] * 128))
+
+    task = dataset.create_data_rows([{
+        DataRow.row_data: image_url,
+        DataRow.custom_metadata: fields
+    }])
+    task.wait_till_done()
+    assert task.status == "FAILED"
+
+
+def test_create_data_rows_with_metadata_missing_value(dataset, image_url):
+    fields = make_metadata_fields()
+    fields.append({"schemaId": "some schema id"})
+
+    with pytest.raises(ValueError) as exc:
+        dataset.create_data_rows([
+            {
+                DataRow.row_data: image_url,
+                DataRow.external_id: "row1",
+                DataRow.custom_metadata: fields
+            },
+        ])
+
+
+def test_create_data_rows_with_metadata_missing_schema_id(dataset, image_url):
+    fields = make_metadata_fields()
+    fields.append({"value": "some value"})
+
+    with pytest.raises(ValueError) as exc:
+        dataset.create_data_rows([
+            {
+                DataRow.row_data: image_url,
+                DataRow.external_id: "row1",
+                DataRow.custom_metadata: fields
+            },
+        ])
+
+
+def test_create_data_rows_with_metadata_wrong_type(dataset, image_url):
+    fields = make_metadata_fields()
+    fields.append("Neither DataRowMetadataField or dict")
+
+    with pytest.raises(ValueError) as exc:
+        task = dataset.create_data_rows([
+            {
+                DataRow.row_data: image_url,
+                DataRow.external_id: "row1",
+                DataRow.custom_metadata: fields
+            },
+        ])
 
 
 def test_data_row_update(dataset, rand_gen, image_url):