fix(content-script-ui): Properly assign and unassign mounted value (#598)

Author: aklinker1

Diff for: src/client/content-scripts/ui/__tests__/index.test.ts

@@ -404,4 +404,62 @@ describe('Content Script UIs', () => {
       });
     });
   });
+
+  describe('mounted value', () => {
+    describe('integrated', () => {
+      it('should set the mounted value based on the onMounted return value', () => {
+        const expected = Symbol();
+
+        const ui = createIntegratedUi(new ContentScriptContext('test'), {
+          position: 'inline',
+          onMount: () => expected,
+        });
+        expect(ui.mounted).toBeUndefined();
+
+        ui.mount();
+        expect(ui.mounted).toBe(expected);
+
+        ui.remove();
+        expect(ui.mounted).toBeUndefined();
+      });
+    });
+
+    describe('iframe', () => {
+      it('should set the mounted value based on the onMounted return value', async () => {
+        const expected = Symbol();
+
+        const ui = createIframeUi(new ContentScriptContext('test'), {
+          page: '',
+          position: 'inline',
+          onMount: () => expected,
+        });
+        expect(ui.mounted).toBeUndefined();
+
+        ui.mount();
+        expect(ui.mounted).toBe(expected);
+
+        ui.remove();
+        expect(ui.mounted).toBeUndefined();
+      });
+    });
+
+    describe('shadow-root', () => {
+      it('should set the mounted value based on the onMounted return value', async () => {
+        const expected = Symbol();
+
+        const ui = await createShadowRootUi(new ContentScriptContext('test'), {
+          name: 'test',
+          position: 'inline',
+          onMount: () => expected,
+        });
+        expect(ui.mounted).toBeUndefined();
+
+        ui.mount();
+        expect(ui.mounted).toBe(expected);
+
+        ui.remove();
+        expect(ui.mounted).toBeUndefined();
+      });
+    });
+  });
 });

Diff for: src/client/content-scripts/ui/index.ts

@@ -35,12 +35,15 @@ export function createIntegratedUi<TMounted>(
   const remove = () => {
     options.onRemove?.(mounted);
     wrapper.remove();
+    mounted = undefined;
   };
 
   ctx.onInvalidated(remove);
 
   return {
-    mounted,
+    get mounted() {
+      return mounted;
+    },
     wrapper,
     mount,
     remove,
@@ -71,12 +74,15 @@ export function createIframeUi<TMounted>(
   const remove = () => {
     options.onRemove?.(mounted);
     wrapper.remove();
+    mounted = undefined;
   };
 
   ctx.onInvalidated(remove);
 
   return {
-    mounted,
+    get mounted() {
+      return mounted;
+    },
     iframe,
     wrapper,
     mount,
@@ -116,7 +122,7 @@ export async function createShadowRootUi<TMounted>(
   });
   shadowHost.setAttribute('data-wxt-shadow-root', '');
 
-  let mounted: TMounted;
+  let mounted: TMounted | undefined;
 
   const mount = () => {
     // Add shadow root element to DOM
@@ -134,6 +140,8 @@ export async function createShadowRootUi<TMounted>(
     // Remove children from uiContainer
     while (uiContainer.lastChild)
       uiContainer.removeChild(uiContainer.lastChild);
+    // Clear mounted value
+    mounted = undefined;
   };
 
   ctx.onInvalidated(remove);
@@ -144,7 +152,9 @@ export async function createShadowRootUi<TMounted>(
     uiContainer,
     mount,
     remove,
-    mounted: mounted!,
+    get mounted() {
+      return mounted;
+    },
   };
 }