From c605c6518b1ceb6077ed4f959bf7adb815c74afd Mon Sep 17 00:00:00 2001
From: Parbhat Puri <parbhatpuri17@gmail.com>
Date: Thu, 2 Apr 2020 19:19:44 +0530
Subject: [PATCH] Fix snippet embed display in editor and frontend

---
 wagtail_draftail_snippet/richtext.py              |  9 ++++-----
 .../js/wagtail_draftail_snippet.js                | 15 +++++----------
 2 files changed, 9 insertions(+), 15 deletions(-)

diff --git a/wagtail_draftail_snippet/richtext.py b/wagtail_draftail_snippet/richtext.py
index ab3891f..e6d8705 100644
--- a/wagtail_draftail_snippet/richtext.py
+++ b/wagtail_draftail_snippet/richtext.py
@@ -144,12 +144,11 @@ class SnippetEmbedElementHandler(AtomicBlockEntityElementHandler):
 
     def create_entity(self, name, attrs, state, contentstate):
         return Entity('SNIPPET-EMBED', 'IMMUTABLE', {
-            "embedtype": "snippet",
             "id": attrs.get("id"),
-            "data-string": attrs.get("string"),
-            "data-edit-link": attrs.get("edit_link"),
-            "data-app-name": attrs.get("app_name"),
-            "data-model-name": attrs.get("model_name"),
+            "string": attrs.get("data-string"),
+            "edit_link": attrs.get("data-edit-link"),
+            "app_name": attrs.get("data-app-name"),
+            "model_name": attrs.get("data-model-name"),
         })
 
 
diff --git a/wagtail_draftail_snippet/static/wagtail_draftail_snippet/js/wagtail_draftail_snippet.js b/wagtail_draftail_snippet/static/wagtail_draftail_snippet/js/wagtail_draftail_snippet.js
index 4303fc4..15af48e 100644
--- a/wagtail_draftail_snippet/static/wagtail_draftail_snippet/js/wagtail_draftail_snippet.js
+++ b/wagtail_draftail_snippet/static/wagtail_draftail_snippet/js/wagtail_draftail_snippet.js
@@ -182,20 +182,15 @@
   };
 
   const SnippetEmbed = props => {
-    const { entity, onRemoveEntity } = props.blockProps;
+    const { entity, onRemoveEntity, entityKey } = props.blockProps;
     const data = entity.getData();
 
-    let icon = React.createElement(window.wagtail.components.Icon, {name: 'media'});
+    let icon = React.createElement(window.wagtail.components.Icon, {name: 'snippet'});
     let label = data.string || '';
 
-    return React.createElement(TooltipEntity, {
-      entityKey: props.entityKey,
-      children: props.children,
-      onEdit: props.onEdit,
-      onRemove: props.onRemove,
-      icon: icon,
-      label: label
-    });
+    return React.createElement("div", {
+      class: "MediaBlock"
+    }, icon, `${label}`);
   };
 
   window.draftail.registerPlugin({