diff --git a/src/mediaembedcommand.js b/src/mediaembedcommand.js
index 94802ce..3ef85b4 100644
--- a/src/mediaembedcommand.js
+++ b/src/mediaembedcommand.js
@@ -9,7 +9,7 @@
 
 import Command from '@ckeditor/ckeditor5-core/src/command';
 import { findOptimalInsertionPosition } from '@ckeditor/ckeditor5-widget/src/utils';
-import { getSelectedMediaElement, insertMedia } from './utils';
+import { getSelectedMediaModelWidget, insertMedia } from './utils';
 
 /**
  * The insert media command.
@@ -31,7 +31,7 @@ export default class MediaEmbedCommand extends Command {
 		const selection = model.document.selection;
 		const schema = model.schema;
 		const position = selection.getFirstPosition();
-		const selectedMedia = getSelectedMediaElement( selection );
+		const selectedMedia = getSelectedMediaModelWidget( selection );
 
 		let parent = position.parent;
 
@@ -55,7 +55,7 @@ export default class MediaEmbedCommand extends Command {
 	execute( url ) {
 		const model = this.editor.model;
 		const selection = model.document.selection;
-		const selectedMedia = getSelectedMediaElement( selection );
+		const selectedMedia = getSelectedMediaModelWidget( selection );
 
 		if ( selectedMedia ) {
 			model.change( writer => {
diff --git a/src/mediaembedtoolbar.js b/src/mediaembedtoolbar.js
index be02931..57b1b25 100644
--- a/src/mediaembedtoolbar.js
+++ b/src/mediaembedtoolbar.js
@@ -8,7 +8,7 @@
 
 import Plugin from '@ckeditor/ckeditor5-core/src/plugin';
 import WidgetToolbarRepository from '@ckeditor/ckeditor5-widget/src/widgettoolbarrepository';
-import { isMediaWidgetSelected } from './utils';
+import { getSelectedMediaViewWidget } from './utils';
 
 /**
  * The media embed toolbar plugin. It creates a toolbar for media embed that shows up when the media element is selected.
@@ -42,7 +42,7 @@ export default class MediaEmbedToolbar extends Plugin {
 
 		widgetToolbarRepository.register( 'mediaEmbed', {
 			items: editor.config.get( 'mediaEmbed.toolbar' ) || [],
-			visibleWhen: isMediaWidgetSelected,
+			getRelatedElement: getSelectedMediaViewWidget,
 		} );
 	}
 }
diff --git a/src/utils.js b/src/utils.js
index 4715d62..f34078b 100644
--- a/src/utils.js
+++ b/src/utils.js
@@ -27,12 +27,28 @@ export function toMediaWidget( viewElement, writer, label ) {
 	return toWidget( viewElement, writer, { label } );
 }
 
-export function isMediaWidgetSelected( viewSelection ) {
-	const viewElement = viewSelection.getSelectedElement();
+/**
+ * Returns a media widget editing view element if one is selected.
+ *
+ * @param {module:engine/view/selection~Selection|module:engine/view/documentselection~DocumentSelection} selection
+ * @returns {module:engine/view/element~Element|null}
+ */
+export function getSelectedMediaViewWidget( selection ) {
+	const viewElement = selection.getSelectedElement();
+
+	if ( viewElement && isMediaWidget( viewElement ) ) {
+		return viewElement;
+	}
 
-	return !!( viewElement && isMediaWidget( viewElement ) );
+	return null;
 }
 
+/**
+ * Checks if a given view element is a media widget.
+ *
+ * @param {module:engine/view/element~Element} viewElement
+ * @returns {Boolean}
+ */
 export function isMediaWidget( viewElement ) {
 	return !!viewElement.getCustomProperty( mediaSymbol ) && isWidget( viewElement );
 }
@@ -78,7 +94,7 @@ export function createMediaFigureElement( writer, registry, url, options ) {
  * @param {module:engine/model/selection~Selection} selection
  * @returns {module:engine/model/element~Element|null}
  */
-export function getSelectedMediaElement( selection ) {
+export function getSelectedMediaModelWidget( selection ) {
 	const selectedElement = selection.getSelectedElement();
 
 	if ( selectedElement && selectedElement.is( 'media' ) ) {
diff --git a/tests/mediaembedtoolbar.js b/tests/mediaembedtoolbar.js
index eaca3b7..92c7f98 100644
--- a/tests/mediaembedtoolbar.js
+++ b/tests/mediaembedtoolbar.js
@@ -35,7 +35,7 @@ describe( 'MediaEmbedToolbar', () => {
 			editor = _editor;
 			model = editor.model;
 			widgetToolbarRepository = editor.plugins.get( 'WidgetToolbarRepository' );
-			toolbar = widgetToolbarRepository._toolbars.get( 'mediaEmbed' ).view;
+			toolbar = widgetToolbarRepository._toolbarDefinitions.get( 'mediaEmbed' ).view;
 			balloon = editor.plugins.get( 'ContextualBalloon' );
 		} );
 	} );
@@ -184,7 +184,7 @@ describe( 'MediaEmbedToolbar - integration with BalloonEditor', () => {
 			editor = _editor;
 			model = editor.model;
 			widgetToolbarRepository = editor.plugins.get( 'WidgetToolbarRepository' );
-			toolbar = widgetToolbarRepository._toolbars.get( 'mediaEmbed' ).view;
+			toolbar = widgetToolbarRepository._toolbarDefinitions.get( 'mediaEmbed' ).view;
 			balloon = editor.plugins.get( 'ContextualBalloon' );
 			balloonToolbar = editor.plugins.get( 'BalloonToolbar' );
 		} );