diff --git a/src/plugins/vis_augmenter/public/saved_augment_vis/utils/helpers.ts b/src/plugins/vis_augmenter/public/saved_augment_vis/utils/helpers.ts
index ff23a229b154..0751f1dd330c 100644
--- a/src/plugins/vis_augmenter/public/saved_augment_vis/utils/helpers.ts
+++ b/src/plugins/vis_augmenter/public/saved_augment_vis/utils/helpers.ts
@@ -4,21 +4,27 @@
  */
 
 import { get } from 'lodash';
-import { getSavedAugmentVisLoader, getUISettings } from '../../services';
 import { ISavedAugmentVis } from '../types';
 import {
   PLUGIN_AUGMENTATION_ENABLE_SETTING,
   PLUGIN_AUGMENTATION_MAX_OBJECTS_SETTING,
 } from '../../../common/constants';
+import { SavedAugmentVisLoader } from '../saved_augment_vis';
+import { getSavedAugmentVisLoader, getUISettings } from '../../services';
+import { IUiSettingsClient } from '../../../../../core/public';
 
 /**
  * Create an augment vis saved object given an object that
  * implements the ISavedAugmentVis interface
  */
-export const createAugmentVisSavedObject = async (AugmentVis: ISavedAugmentVis): Promise<any> => {
-  const loader = getSavedAugmentVisLoader();
-  const config = getUISettings();
-
+export const createAugmentVisSavedObject = async (
+  AugmentVis: ISavedAugmentVis,
+  savedObjLoader?: SavedAugmentVisLoader,
+  uiSettings?: IUiSettingsClient
+): Promise<any> => {
+  // Using optional services provided, or the built-in services from this plugin
+  const loader = savedObjLoader !== undefined ? savedObjLoader : getSavedAugmentVisLoader();
+  const config = uiSettings !== undefined ? uiSettings : getUISettings();
   const isAugmentationEnabled = config.get(PLUGIN_AUGMENTATION_ENABLE_SETTING);
   if (!isAugmentationEnabled) {
     throw new Error(
diff --git a/src/plugins/vis_augmenter/public/utils/utils.ts b/src/plugins/vis_augmenter/public/utils/utils.ts
index ead35fa4b33a..9d683ff962c1 100644
--- a/src/plugins/vis_augmenter/public/utils/utils.ts
+++ b/src/plugins/vis_augmenter/public/utils/utils.ts
@@ -20,6 +20,7 @@ import {
 } from '../';
 import { PLUGIN_AUGMENTATION_ENABLE_SETTING } from '../../common/constants';
 import { getUISettings } from '../services';
+import { IUiSettingsClient } from '../../../../core/public';
 
 export const isEligibleForVisLayers = (vis: Vis): boolean => {
   // Only support date histogram and ensure there is only 1 x-axis and it has to be on the bottom.
@@ -53,9 +54,11 @@ export const isEligibleForVisLayers = (vis: Vis): boolean => {
  */
 export const getAugmentVisSavedObjs = async (
   visId: string | undefined,
-  loader: SavedAugmentVisLoader | undefined
+  loader: SavedAugmentVisLoader | undefined,
+  uiSettings?: IUiSettingsClient | undefined
 ): Promise<ISavedAugmentVis[]> => {
-  const config = getUISettings();
+  // Using optional services provided, or the built-in services from this plugin
+  const config = uiSettings !== undefined ? uiSettings : getUISettings();
   const isAugmentationEnabled = config.get(PLUGIN_AUGMENTATION_ENABLE_SETTING);
   if (!isAugmentationEnabled) {
     throw new Error(