1 files changed, 13 insertions, 6 deletions
diff --git a/app/assets/javascripts/image_diff/helpers/utils_helper.js b/app/assets/javascripts/image_diff/helpers/utils_helper.js
index d0cd42df073..4015f2db187 100644
--- a/app/assets/javascripts/image_diff/helpers/utils_helper.js
+++ b/app/assets/javascripts/image_diff/helpers/utils_helper.js
@@ -1,6 +1,7 @@
 import ImageBadge from '../image_badge';
 import ImageDiff from '../image_diff';
 import ReplacedImageDiff from '../replaced_image_diff';
+import '../../commit/image_file';
 
 export function resizeCoordinatesToImageElement(imageEl, meta) {
   const { x, y, width, height } = meta;
@@ -76,7 +77,7 @@ export function getTargetSelection(event) {
   };
 }
 
-export function initImageDiff(file, canCreateNote, renderCommentBadge) {
+export function initImageDiff(fileEl, canCreateNote, renderCommentBadge) {
   const options = {
     canCreateNote,
     renderCommentBadge,
@@ -84,13 +85,19 @@ export function initImageDiff(file, canCreateNote, renderCommentBadge) {
 
   // ImageFile needs to be invoked before initImageDiff so that badges
   // can mount to the correct location
-  new gl.ImageFile(file); // eslint-disable-line no-new
+  new gl.ImageFile(fileEl); // eslint-disable-line no-new
 
-  if (file.querySelector('.diff-file .js-single-image')) {
-    const imageDiff = new ImageDiff(file, options);
+  if (fileEl.querySelector('.diff-file .js-single-image')) {
+    const imageDiff = new ImageDiff(fileEl, options);
     imageDiff.init();
-  } else if (file.querySelector('.diff-file .js-replaced-image')) {
-    const replacedImageDiff = new ReplacedImageDiff(file, options);
+
+    return imageDiff;
+  } else if (fileEl.querySelector('.diff-file .js-replaced-image')) {
+    const replacedImageDiff = new ReplacedImageDiff(fileEl, options);
     replacedImageDiff.init();
+
+    return replacedImageDiff;
   }
+
+  return null;
 }