diff --git a/x-pack/plugins/lens/public/app_plugin/app.test.tsx b/x-pack/plugins/lens/public/app_plugin/app.test.tsx
index ee08d86061a3f..215dfe9b4ec9b 100644
--- a/x-pack/plugins/lens/public/app_plugin/app.test.tsx
+++ b/x-pack/plugins/lens/public/app_plugin/app.test.tsx
@@ -38,18 +38,6 @@ import faker from 'faker';
 import { screen, waitFor } from '@testing-library/react';
 import userEvent from '@testing-library/user-event';
 import { VisualizeEditorContext } from '../types';
-import { pluginServices as presentationUtilPluginServices } from '@kbn/presentation-util-plugin/public/services';
-
-presentationUtilPluginServices.getHooks = jest.fn().mockReturnValue({
-  capabilities: {
-    useService: () => ({
-      canAccessDashboards: () => true,
-      canCreateNewDashboards: () => true,
-      canSaveVisualizations: () => true,
-      canSetAdvancedSettings: () => true,
-    }),
-  },
-});
 
 jest.mock('../persistence/saved_objects_utils/check_for_duplicate_title', () => ({
   checkForDuplicateTitle: jest.fn(),
diff --git a/x-pack/plugins/lens/public/app_plugin/mounter.tsx b/x-pack/plugins/lens/public/app_plugin/mounter.tsx
index 6cb3fa37993fb..93c5543372596 100644
--- a/x-pack/plugins/lens/public/app_plugin/mounter.tsx
+++ b/x-pack/plugins/lens/public/app_plugin/mounter.tsx
@@ -5,7 +5,7 @@
  * 2.0.
  */
 
-import React, { useCallback, useEffect, useState, useMemo, PropsWithChildren, FC } from 'react';
+import React, { useCallback, useEffect, useState, useMemo } from 'react';
 import { AppMountParameters, CoreSetup, CoreStart } from '@kbn/core/public';
 import { FormattedMessage } from '@kbn/i18n-react';
 import { RouteComponentProps } from 'react-router-dom';
@@ -143,7 +143,6 @@ export async function mountApp(
   mountProps: {
     createEditorFrame: EditorFrameStart['createInstance'];
     attributeService: LensAttributesService;
-    getPresentationUtilContext: () => FC<PropsWithChildren<{}>>;
     topNavMenuEntryGenerators: LensTopNavMenuEntryGenerator[];
     locator?: LensAppLocator;
   }
diff --git a/x-pack/plugins/lens/public/react_embeddable/inline_editing/state_management.tsx b/x-pack/plugins/lens/public/react_embeddable/inline_editing/state_management.tsx
index cdd2f60edda77..2a4f1f48fd0dc 100644
--- a/x-pack/plugins/lens/public/react_embeddable/inline_editing/state_management.tsx
+++ b/x-pack/plugins/lens/public/react_embeddable/inline_editing/state_management.tsx
@@ -25,7 +25,7 @@ export function getStateManagementForInlineEditing(
   const updatePanelState = (
     datasourceState: unknown,
     visualizationState: unknown,
-    visualizationType: string
+    visualizationType?: string
   ) => {
     const viz = getAttributes();
     const datasourceStates: DatasourceStates = {
@@ -37,7 +37,7 @@ export function getStateManagementForInlineEditing(
     const newViz = mergeToNewDoc(
       viz,
       {
-        activeId: visualizationType,
+        activeId: visualizationType || viz.visualizationType,
         state: visualizationState,
       },
       datasourceStates,