diff --git a/tests/integration/annotation_import/conftest.py b/tests/integration/annotation_import/conftest.py
index 56724d2ae..83215a6ba 100644
--- a/tests/integration/annotation_import/conftest.py
+++ b/tests/integration/annotation_import/conftest.py
@@ -7,7 +7,7 @@
 
 from typing import Type
 from labelbox.schema.labeling_frontend import LabelingFrontend
-from labelbox.schema.annotation_import import MALPredictionImport, AnnotationImportState
+from labelbox.schema.annotation_import import LabelImport, MALPredictionImport, AnnotationImportState
 
 
 @pytest.fixture
@@ -327,18 +327,14 @@ def model_run(rand_gen, model):
 
 @pytest.fixture
 def model_run_with_model_run_data_rows(client, configured_project,
-                                       annotation_submit_fn,
                                        model_run_predictions, model_run):
     configured_project.enable_model_assisted_labeling()
 
-    upload_task = MALPredictionImport.create_from_objects(
-        client, configured_project.uid, f'mal-import-{uuid.uuid4()}',
+    upload_task = LabelImport.create_from_objects(
+        client, configured_project.uid, f"label-import-{uuid.uuid4()}",
         model_run_predictions)
     upload_task.wait_until_done()
-    label_ids = []
-    for data_row_id in {x['dataRow']['id'] for x in model_run_predictions}:
-        label_ids.append(
-            annotation_submit_fn(configured_project.uid, data_row_id))
+    label_ids = [label.uid for label in configured_project.labels()]
     model_run.upsert_labels(label_ids)
     time.sleep(3)
     yield model_run
diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py
index f67506116..03c743f88 100644
--- a/tests/integration/conftest.py
+++ b/tests/integration/conftest.py
@@ -21,6 +21,7 @@
 from labelbox.schema.invite import Invite
 from labelbox.schema.user import User
 from labelbox import OntologyBuilder, Tool
+from labelbox.schema.annotation_import import LabelImport
 
 IMG_URL = "https://picsum.photos/200/300"
 
@@ -298,37 +299,7 @@ def configured_project(project, client, rand_gen, image_url):
 
 
 @pytest.fixture
-def annotation_submit_fn(client):
-
-    def submit(project_id, data_row_id):
-        feature_result = client.execute(
-            """query featuresPyApi ($project_id : ID!, $datarow_id: ID!
-            ) {project(where: { id: $project_id }) {
-                    featuresForDataRow(where: {dataRow: { id: $datarow_id }}) {id}}}
-            """, {
-                "project_id": project_id,
-                "datarow_id": data_row_id
-            })
-        features = feature_result['project']['featuresForDataRow']
-        feature_ids = [feature['id'] for feature in features]
-        res = client.execute(
-            """mutation createLabelPyApi ($project_id : ID!,$datarow_id: ID!,$feature_ids: [ID!]!,$time_seconds : Float!) {
-                createLabelFromFeatures(data: {dataRow: { id: $datarow_id },project: { id: $project_id },
-                    featureIds: $feature_ids,secondsSpent: $time_seconds}) {id}}""",
-            {
-                "project_id": project_id,
-                "datarow_id": data_row_id,
-                "feature_ids": feature_ids,
-                "time_seconds": 10
-            })
-        return res['createLabelFromFeatures']['id']
-
-    return submit
-
-
-@pytest.fixture
-def configured_project_with_label(client, rand_gen, annotation_submit_fn,
-                                  image_url):
+def configured_project_with_label(client, rand_gen, image_url):
     project = client.create_project(name=rand_gen(str))
     dataset = client.create_dataset(name=rand_gen(str), projects=project)
     data_row = dataset.create_data_row(row_data=image_url)
@@ -355,11 +326,10 @@ def configured_project_with_label(client, rand_gen, annotation_submit_fn,
             "width": 50
         }
     }]
-    upload_task = MALPredictionImport.create_from_objects(
-        client, project.uid, f'mal-import-{uuid.uuid4()}', predictions)
+    upload_task = LabelImport.create_from_objects(
+        client, project.uid, f'label-import-{uuid.uuid4()}', predictions)
     upload_task.wait_until_done()
-    labels = annotation_submit_fn(project.uid, data_row.uid)
-    time.sleep(3)
-    yield [project, labels]
+    label = next(project.labels()).uid
+    yield [project, label]
     dataset.delete()
     project.delete()