Improve preview insertion logic (#191)

src/preview.ts

@@ -1,5 +1,4 @@
 import {formatCause} from '@croct/sdk/error';
-import {uuid4} from '@croct/sdk/uuid';
 import {Logger} from './sdk';
 import {Plugin, PluginFactory} from './plugin';
 import {Token, TokenStore} from './sdk/token';
@@ -25,7 +24,7 @@ export class PreviewPlugin implements Plugin {
 
     private readonly logger: Logger;
 
-    private readonly widgetId = `croct-preview:${uuid4()}`;
+    private cleanUp: Array<() => void> = [];
 
     public constructor(configuration: Configuration) {
         this.tokenStore = configuration.tokenStore;
@@ -55,7 +54,9 @@ export class PreviewPlugin implements Plugin {
     }
 
     public disable(): void {
-        document.getElementById(this.widgetId)?.remove();
+        const callbacks = this.cleanUp.splice(0);
+
+        callbacks.forEach(cleanUp => cleanUp());
     }
 
     private updateToken(data: string): void {
@@ -77,6 +78,8 @@ export class PreviewPlugin implements Plugin {
                 token = null;
             }
 
+            this.logger.debug('Preview token updated.');
+
             this.tokenStore.setToken(token);
         } catch (error) {
             this.tokenStore.setToken(null);
@@ -118,7 +121,7 @@ export class PreviewPlugin implements Plugin {
     private insertWidget(url: string): void {
         const widget = this.createWidget(url);
 
-        window.addEventListener('message', event => {
+        const onMessage = (event: MessageEvent): void => {
             if (event.origin !== PREVIEW_WIDGET_ORIGIN) {
                 return;
             }
@@ -142,13 +145,42 @@ export class PreviewPlugin implements Plugin {
 
                     break;
             }
-        });
+        };
+
+        window.addEventListener('message', onMessage);
+
+        this.cleanUp.push(() => window.removeEventListener('message', onMessage));
+
+        const insert = (): void => {
+            const container = document.body;
+
+            container.prepend(widget);
+
+            this.cleanUp.push(() => container.removeChild(widget));
+
+            this.logger.debug('Preview widget mounted.');
+
+            // Reinsert the widget if it is removed from the DOM (e.g. by a framework)
+            const observer = new MutationObserver(() => {
+                if (!container.contains(widget)) {
+                    container.prepend(widget);
+
+                    this.logger.debug('Preview widget removed from DOM, reinserting.');
+                }
+            });
+
+            observer.observe(container, {childList: true});
 
-        const insert = (): void => document.body?.prepend(widget);
+            // Make sure the observer is disconnected before the widget is removed.
+            // Otherwise, the observer may reinsert the widget after it is removed.
+            this.cleanUp.unshift(() => observer.disconnect());
+        };
 
-        if (document.readyState === 'complete') {
+        if (document.readyState !== 'loading') {
             insert();
         } else {
+            this.logger.debug('Waiting for document to be ready.');
+            this.cleanUp.push(() => window.removeEventListener('DOMContentLoaded', insert));
             window.addEventListener('DOMContentLoaded', insert);
         }
     }
@@ -166,7 +198,6 @@ export class PreviewPlugin implements Plugin {
     private createWidget(url: string): HTMLIFrameElement {
         const widget = document.createElement('iframe');
 
-        widget.setAttribute('id', this.widgetId);
         widget.setAttribute('src', url);
         widget.setAttribute('sandbox', 'allow-scripts allow-same-origin');
 

test/preview.test.ts

@@ -39,10 +39,19 @@ describe('A preview plugin factory', () => {
 describe('A Preview plugin', () => {
     beforeEach(() => {
         window.history.replaceState({}, 'Home page', 'http://localhost');
-        document.body.innerHTML = '';
+
+        document.getElementsByTagName('html')[0].innerHTML = '';
 
         const {location} = window;
 
+        jest.spyOn(window, 'removeEventListener');
+        jest.spyOn(window, 'addEventListener');
+
+        Object.defineProperty(document, 'readyState', {
+            writable: true,
+            value: 'complete',
+        });
+
         Object.defineProperty(
             window,
             'location',
@@ -62,6 +71,12 @@ describe('A Preview plugin', () => {
     });
 
     afterEach(() => {
+        const {calls} = jest.mocked(window.addEventListener).mock;
+
+        for (const [event, listener] of calls) {
+            window.removeEventListener(event, listener);
+        }
+
         jest.resetAllMocks();
     });
 
@@ -267,14 +282,13 @@ describe('A Preview plugin', () => {
 
         expect(document.body.querySelector('iframe')).toBe(null);
 
-        Object.defineProperty(document, 'readyState', {
-            writable: true,
-            value: 'complete',
-        });
+        expect(configuration.logger.debug).toHaveBeenCalledWith('Waiting for document to be ready.');
 
         window.dispatchEvent(new Event('DOMContentLoaded'));
 
         expect(document.body.querySelector('iframe')).not.toBe(null);
+
+        expect(configuration.logger.debug).toHaveBeenCalledWith('Preview widget mounted.');
     });
 
     it('should ignore messages from unknown origins', () => {
@@ -358,18 +372,76 @@ describe('A Preview plugin', () => {
         expect(window.location.replace).toHaveBeenCalledWith('http://localhost/?croct-preview=exit');
     });
 
-    it('should remove widget', () => {
+    it('should remount the widget when removed from the DOM', () => {
+        const observeSpy = jest.spyOn(window, 'MutationObserver')
+            .mockImplementation(
+                () => ({
+                    observe: jest.fn(),
+                    disconnect: jest.fn(),
+                    takeRecords: jest.fn(),
+                }),
+            );
+
         const plugin = new PreviewPlugin(configuration);
 
         configuration.tokenStore.setToken(token);
 
         plugin.enable();
 
+        const widget = document.body.querySelector('iframe') as HTMLIFrameElement;
+
+        expect(widget).not.toBe(null);
+
+        document.body.removeChild(widget);
+
+        expect(document.body.querySelector('iframe')).toBe(null);
+
+        const callback: MutationCallback = observeSpy.mock.calls[0][0];
+        const instance: MutationObserver = observeSpy.mock.instances[0];
+
+        callback([], instance);
+
         expect(document.body.querySelector('iframe')).not
             .toBe(null);
 
+        expect(configuration.logger.debug).toHaveBeenCalledWith('Preview widget removed from DOM, reinserting.');
+    });
+
+    it('should remove widget', () => {
+        const disconnectSpy = jest.fn();
+
+        jest.spyOn(window, 'MutationObserver')
+            .mockImplementation(
+                () => ({
+                    observe: jest.fn(),
+                    disconnect: disconnectSpy,
+                    takeRecords: jest.fn(),
+                }),
+            );
+
+        Object.defineProperty(document, 'readyState', {
+            writable: true,
+            value: 'loading',
+        });
+
+        const plugin = new PreviewPlugin(configuration);
+
+        configuration.tokenStore.setToken(token);
+
+        expect(document.body.children.length).toBe(0);
+
+        plugin.enable();
+
+        window.dispatchEvent(new Event('DOMContentLoaded'));
+
+        expect(document.body.children.length).toBe(1);
+
         plugin.disable();
 
-        expect(document.body.querySelector('iframe')).toBe(null);
+        expect(window.removeEventListener).toHaveBeenCalledWith('message', expect.any(Function));
+        expect(window.removeEventListener).toHaveBeenCalledWith('DOMContentLoaded', expect.any(Function));
+        expect(disconnectSpy).toHaveBeenCalled();
+
+        expect(document.body.children.length).toBe(0);
     });
 });