From ca86446c91045576c3062a773db11f9dc4dc7177 Mon Sep 17 00:00:00 2001
From: "Julia Roldi (from Dev Box)" <juliaroldi@microsoft.com>
Date: Tue, 30 Apr 2024 14:32:36 -0300
Subject: [PATCH] fixes

---
 .../setDOMSelection/setDOMSelection.ts        | 28 ++++++++-----------
 .../optimizers/removeUnnecessarySpan.ts       |  5 ++--
 2 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/packages/roosterjs-content-model-core/lib/coreApi/setDOMSelection/setDOMSelection.ts b/packages/roosterjs-content-model-core/lib/coreApi/setDOMSelection/setDOMSelection.ts
index 158da4db215..a89fee21134 100644
--- a/packages/roosterjs-content-model-core/lib/coreApi/setDOMSelection/setDOMSelection.ts
+++ b/packages/roosterjs-content-model-core/lib/coreApi/setDOMSelection/setDOMSelection.ts
@@ -7,6 +7,7 @@ import {
     isNodeOfType,
     parseTableCells,
     toArray,
+    wrap,
 } from 'roosterjs-content-model-dom';
 import type {
     ParsedTable,
@@ -44,14 +45,12 @@ export const setDOMSelection: SetDOMSelection = (core, selection, skipSelectionC
     try {
         switch (selection?.type) {
             case 'image':
-                const image = selection.image;
+                const image = ensureImageHasSpanParent(selection.image);
 
-                const spanImage = ensureImageHasSpanParent(image);
-                if (!spanImage) {
-                    return;
-                }
-
-                core.selection.selection = selection;
+                core.selection.selection = {
+                    type: 'image',
+                    image,
+                };
                 core.api.setEditorStyle(
                     core,
                     DOM_SELECTION_CSS_KEY,
@@ -242,22 +241,19 @@ function setRangeSelection(doc: Document, element: HTMLElement | undefined, coll
     }
 }
 
-function ensureImageHasSpanParent(image: HTMLImageElement) {
+function ensureImageHasSpanParent(image: HTMLImageElement): HTMLImageElement {
     const parent = image.parentElement;
 
     if (
         parent &&
         isNodeOfType(parent, 'ELEMENT_NODE') &&
         isElementOfType(parent, 'span') &&
-        parent.firstElementChild == image &&
-        parent.lastElementChild == image &&
-        !parent.textContent
+        parent.firstChild == image &&
+        parent.lastChild == image
     ) {
-        return parent;
+        return image;
     }
 
-    const span = image.ownerDocument.createElement('span');
-    span.appendChild(image);
-    parent?.appendChild(span);
-    return parent;
+    wrap(image.ownerDocument, image, 'span');
+    return image;
 }
diff --git a/packages/roosterjs-content-model-dom/lib/modelToDom/optimizers/removeUnnecessarySpan.ts b/packages/roosterjs-content-model-dom/lib/modelToDom/optimizers/removeUnnecessarySpan.ts
index 891846aeecc..3b0ccb11b88 100644
--- a/packages/roosterjs-content-model-dom/lib/modelToDom/optimizers/removeUnnecessarySpan.ts
+++ b/packages/roosterjs-content-model-dom/lib/modelToDom/optimizers/removeUnnecessarySpan.ts
@@ -30,7 +30,8 @@ export function removeUnnecessarySpan(root: Node) {
 
 const isImageSpan = (child: HTMLElement) => {
     return (
-        child?.firstElementChild?.tagName == 'IMG' &&
-        child.firstElementChild == child.lastElementChild
+        isNodeOfType(child.firstChild, 'ELEMENT_NODE') &&
+        child.firstChild.tagName == 'IMG' &&
+        child.firstChild == child.lastChild
     );
 };