Skip to content

Commit 268ee7b

Browse files
Add test
1 parent a4af25f commit 268ee7b

File tree

2 files changed

+33
-20
lines changed

2 files changed

+33
-20
lines changed

labelbox/schema/model_run.py

Lines changed: 19 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -89,22 +89,6 @@ def delete_model_run(self):
8989
ids_param: str(self.uid)
9090
})
9191

92-
93-
class AnnotationGroup(DbObject):
94-
label_id = Field.String("label_id")
95-
model_run_id = Field.String("model_run_id")
96-
data_row = Relationship.ToOne("DataRow", False, cache=True)
97-
98-
def __init__(self, client, model_id, *args, **kwargs):
99-
super().__init__(client, *args, **kwargs)
100-
self.model_id = model_id
101-
102-
@property
103-
def url(self):
104-
app_url = self.client.app_url
105-
endpoint = f"{app_url}/models/{self.model_id}/{self.model_run_id}/AllDatarowsSlice/{self.uid}?view=carousel"
106-
return endpoint
107-
10892
def delete_annotation_groups(self, data_row_ids):
10993
""" Deletes annotation groups by data row ids for a model run.
11094
@@ -115,11 +99,28 @@ def delete_annotation_groups(self, data_row_ids):
11599
"""
116100
model_run_id_param = "modelRunId"
117101
data_row_ids_param = "dataRowIds"
118-
query_str = """mutation DeleteModelRunDataRowsPyApi($%s: ID!, $%s: [ID!]! {
102+
query_str = """mutation DeleteModelRunDataRowsPyApi($%s: ID!, $%s: [ID!]!) {
119103
deleteModelRunDataRows(where: {modelRunId: $%s, dataRowIds: $%s})}""" % (
120104
model_run_id_param, data_row_ids_param, model_run_id_param, data_row_ids_param
121105
)
122106
self.client.execute(query_str, {
123-
model_run_id_param: self.model_run_id,
107+
model_run_id_param: self.uid,
124108
data_row_ids_param: data_row_ids
125109
})
110+
111+
112+
class AnnotationGroup(DbObject):
113+
label_id = Field.String("label_id")
114+
model_run_id = Field.String("model_run_id")
115+
data_row = Relationship.ToOne("DataRow", False, cache=True)
116+
117+
def __init__(self, client, model_id, *args, **kwargs):
118+
super().__init__(client, *args, **kwargs)
119+
self.model_id = model_id
120+
121+
@property
122+
def url(self):
123+
app_url = self.client.app_url
124+
endpoint = f"{app_url}/models/{self.model_id}/{self.model_run_id}/AllDatarowsSlice/{self.uid}?view=carousel"
125+
return endpoint
126+

tests/integration/mal_and_mea/test_model_run.py

Lines changed: 14 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,5 +40,17 @@ def test_model_run_delete(client, model_run):
4040

4141

4242
def test_model_run_delete(client, model_run_annotation_groups):
43-
# TODO
44-
pass
43+
models = list(client.get_models())
44+
model = models[0]
45+
model_runs = list(model.model_runs())
46+
model_run = model_runs[0]
47+
48+
before = list(model_run.annotation_groups())
49+
annotation_group = before[0]
50+
51+
data_row_id = annotation_group.data_row().uid
52+
model_run.delete_annotation_groups(data_row_ids=[data_row_id])
53+
54+
after = list(model_run.annotation_groups())
55+
56+
assert len(before) == len(after) + 1

0 commit comments

Comments
 (0)