diff --git a/test/integration/machine-learning.spec.ts b/test/integration/machine-learning.spec.ts
index 9a358bf0e3..bed3980a49 100644
--- a/test/integration/machine-learning.spec.ts
+++ b/test/integration/machine-learning.spec.ts
@@ -244,11 +244,11 @@ describe('admin.machineLearning', () => {
       return uploadModelToGcs('model1.tflite', 'valid_model.tflite')
         .then((fileName: string) => {
           modelOptions.tfliteModel!.gcsTfliteUri = fileName;
-          createTemporaryModel(modelOptions)
+          return createTemporaryModel(modelOptions)
             .then((createdModel) => {
               expect(createdModel.validationError).to.be.empty;
               expect(createdModel.published).to.be.false;
-              admin.machineLearning().publishModel(createdModel.modelId)
+              return admin.machineLearning().publishModel(createdModel.modelId)
                 .then((publishedModel) => {
                   expect(publishedModel.published).to.be.true;
                 });
@@ -279,14 +279,14 @@ describe('admin.machineLearning', () => {
       return uploadModelToGcs('model1.tflite', 'valid_model.tflite')
         .then((fileName: string) => {
           modelOptions.tfliteModel!.gcsTfliteUri = fileName;
-          createTemporaryModel(modelOptions)
+          return createTemporaryModel(modelOptions)
             .then((createdModel) => {
               expect(createdModel.validationError).to.be.empty;
               expect(createdModel.published).to.be.false;
-              admin.machineLearning().publishModel(createdModel.modelId)
+              return admin.machineLearning().publishModel(createdModel.modelId)
                 .then((publishedModel) => {
                   expect(publishedModel.published).to.be.true;
-                  admin.machineLearning().unpublishModel(publishedModel.modelId)
+                  return admin.machineLearning().unpublishModel(publishedModel.modelId)
                     .then((unpublishedModel) => {
                       expect(unpublishedModel.published).to.be.false;
                     });