Skip to content

Commit 00bee9a

Browse files
author
Matt Sokoloff
committed
fix metadata tests
1 parent a1bdde0 commit 00bee9a

File tree

1 file changed

+20
-12
lines changed

1 file changed

+20
-12
lines changed

tests/integration/test_data_row_metadata.py

Lines changed: 20 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ def big_dataset(dataset: Dataset, image_url):
2727
"row_data": image_url,
2828
"external_id": "my-image"
2929
},
30-
] * 500)
30+
] * 100)
3131
task.wait_till_done()
3232

3333
yield dataset
@@ -105,24 +105,28 @@ def test_bulk_delete_datarow_metadata(datarow, mdo):
105105

106106
def test_bulk_partial_delete_datarow_metadata(datarow, mdo):
107107
"""Delete a single from metadata"""
108-
assert not len(datarow.metadata["fields"])
108+
n_fields = len(datarow.metadata["fields"])
109109

110110
metadata = make_metadata(datarow.uid)
111111
mdo.bulk_upsert([metadata])
112112

113-
assert len(datarow.metadata["fields"])
113+
assert len(datarow.metadata["fields"]) == (n_fields + 5)
114114

115115
mdo.bulk_delete([
116116
DeleteDataRowMetadata(data_row_id=datarow.uid, fields=[TEXT_SCHEMA_ID])
117117
])
118118

119-
assert len(datarow.metadata["fields"]) == 4
119+
assert len(datarow.metadata["fields"]) == (n_fields + 4)
120120

121121

122122
@pytest.mark.slow
123123
def test_large_bulk_delete_datarow_metadata(big_dataset, mdo):
124124
metadata = []
125-
for dr in big_dataset.export_data_rows():
125+
n_fields_start = 0
126+
for idx, dr in enumerate(big_dataset.export_data_rows()):
127+
if idx == 0:
128+
n_fields_start = len(dr.metadata["fields"])
129+
126130
metadata.append(
127131
DataRowMetadata(data_row_id=dr.uid,
128132
fields=[
@@ -148,24 +152,26 @@ def test_large_bulk_delete_datarow_metadata(big_dataset, mdo):
148152
errors = mdo.bulk_delete(deletes)
149153
assert len(errors) == 0
150154
for dr in big_dataset.export_data_rows():
151-
assert len(dr.metadata["fields"]) == 1
155+
assert len(dr.metadata["fields"]) == 1 + n_fields_start
152156
break
153157

154158

155159
def test_bulk_delete_datarow_enum_metadata(datarow: DataRow, mdo):
156160
"""test bulk deletes for non non fields"""
157-
assert not len(datarow.metadata["fields"])
161+
n_fields = len(datarow.metadata["fields"])
158162
metadata = make_metadata(datarow.uid)
159163
metadata.fields = [
160164
m for m in metadata.fields if m.schema_id == SPLIT_SCHEMA_ID
161165
]
162166
mdo.bulk_upsert([metadata])
163-
assert len(datarow.metadata["fields"])
167+
assert len(datarow.metadata["fields"]) == len(
168+
set([x.schema_id for x in metadata.fields] +
169+
[x['schema_id'] for x in datarow.metadata["fields"]]))
164170

165171
mdo.bulk_delete([
166172
DeleteDataRowMetadata(data_row_id=datarow.uid, fields=[SPLIT_SCHEMA_ID])
167173
])
168-
assert not len(datarow.metadata["fields"])
174+
assert len(datarow.metadata["fields"]) == n_fields
169175

170176

171177
def test_raise_enum_upsert_schema_error(datarow, mdo):
@@ -193,7 +199,6 @@ def test_upsert_non_existent_schema_id(datarow, mdo):
193199

194200

195201
def test_delete_non_existent_schema_id(datarow, mdo):
196-
assert not len(datarow.metadata["fields"])
197202
mdo.bulk_delete([
198203
DeleteDataRowMetadata(data_row_id=datarow.uid,
199204
fields=[EMBEDDING_SCHEMA_ID])
@@ -204,15 +209,18 @@ def test_delete_non_existent_schema_id(datarow, mdo):
204209
@pytest.mark.slow
205210
def test_large_bulk_delete_non_existent_schema_id(big_dataset, mdo):
206211
deletes = []
207-
for dr in big_dataset.export_data_rows():
212+
n_fields_start = 0
213+
for idx, dr in enumerate(big_dataset.export_data_rows()):
214+
if idx == 0:
215+
n_fields_start = len(dr.metadata["fields"])
208216
deletes.append(
209217
DeleteDataRowMetadata(data_row_id=dr.uid,
210218
fields=[EMBEDDING_SCHEMA_ID]))
211219
errors = mdo.bulk_delete(deletes)
212220
assert len(errors) == 0
213221

214222
for dr in big_dataset.export_data_rows():
215-
assert not len(dr.metadata["fields"])
223+
assert len(dr.metadata["fields"]) == n_fields_start
216224
break
217225

218226

0 commit comments

Comments
 (0)