fix(runtime): handle lazy-instance promises for connected & disconnected callbacks (#4072)

* Adds a two unit tests highlighting #4070

* Adds proposed fix for #4070

* Run prettier

* Add equivalent waiting for $ in connected callback

* Apply suggested SNC changes

Co-authored-by: Tanner Reits <47483144+tanner-reits@users.noreply.github.com>

* Apply suggested SNC changes

Co-authored-by: Tanner Reits <47483144+tanner-reits@users.noreply.github.com>

---------

Co-authored-by: Tanner Reits <47483144+tanner-reits@users.noreply.github.com>

Author: joewoodhouse

src/runtime/connected-callback.ts

@@ -106,7 +106,11 @@ export const connectedCallback = (elm: d.HostElement) => {
       addHostEventListeners(elm, hostRef, cmpMeta.$listeners$, false);
 
       // fire off connectedCallback() on component instance
-      fireConnectedCallback(hostRef.$lazyInstance$);
+      if (hostRef?.$lazyInstance$) {
+        fireConnectedCallback(hostRef.$lazyInstance$);
+      } else if (hostRef?.$onReadyPromise$) {
+        hostRef.$onReadyPromise$.then(() => fireConnectedCallback(hostRef.$lazyInstance$));
+      }
     }
 
     endConnected();

src/runtime/disconnected-callback.ts

@@ -5,10 +5,18 @@ import type * as d from '../declarations';
 import { PLATFORM_FLAGS } from './runtime-constants';
 import { safeCall } from './update-component';
 
-export const disconnectedCallback = (elm: d.HostElement) => {
+const disconnectInstance = (instance: any) => {
+  if (BUILD.lazyLoad && BUILD.disconnectedCallback) {
+    safeCall(instance, 'disconnectedCallback');
+  }
+  if (BUILD.cmpDidUnload) {
+    safeCall(instance, 'componentDidUnload');
+  }
+};
+
+export const disconnectedCallback = async (elm: d.HostElement) => {
   if ((plt.$flags$ & PLATFORM_FLAGS.isTmpDisconnected) === 0) {
     const hostRef = getHostRef(elm);
-    const instance: any = BUILD.lazyLoad ? hostRef.$lazyInstance$ : elm;
 
     if (BUILD.hostListener) {
       if (hostRef.$rmListeners$) {
@@ -17,11 +25,12 @@ export const disconnectedCallback = (elm: d.HostElement) => {
       }
     }
 
-    if (BUILD.lazyLoad && BUILD.disconnectedCallback) {
-      safeCall(instance, 'disconnectedCallback');
-    }
-    if (BUILD.cmpDidUnload) {
-      safeCall(instance, 'componentDidUnload');
+    if (!BUILD.lazyLoad) {
+      disconnectInstance(elm);
+    } else if (hostRef?.$lazyInstance$) {
+      disconnectInstance(hostRef.$lazyInstance$);
+    } else if (hostRef?.$onReadyPromise$) {
+      hostRef.$onReadyPromise$.then(() => disconnectInstance(hostRef.$lazyInstance$));
     }
   }
 };

src/runtime/test/lifecycle-sync.spec.tsx

@@ -346,4 +346,78 @@ describe('lifecycle sync', () => {
       </cmp-root>
     `);
   });
+
+  it('call disconnectedCallback even if the element is immediately removed', async () => {
+    let connected = 0;
+    let disconnected = 0;
+
+    @Component({ tag: 'cmp-a' })
+    class CmpA {
+      connectedCallback() {
+        connected++;
+      }
+
+      disconnectedCallback() {
+        disconnected++;
+      }
+
+      render() {
+        return <Host></Host>;
+      }
+    }
+
+    const { doc, waitForChanges } = await newSpecPage({
+      components: [CmpA],
+    });
+
+    const a1 = doc.createElement('cmp-a');
+    doc.body.appendChild(a1);
+    a1.remove();
+
+    await waitForChanges();
+
+    expect(connected).toEqual(1);
+    expect(disconnected).toEqual(1);
+  });
+
+  it('calls disconnect and connect when an element is moved in the DOM', async () => {
+    let connected = 0;
+    let disconnected = 0;
+
+    @Component({ tag: 'cmp-a' })
+    class CmpA {
+      connectedCallback() {
+        connected++;
+      }
+
+      disconnectedCallback() {
+        disconnected++;
+      }
+
+      render() {
+        return <Host></Host>;
+      }
+    }
+
+    const { doc, waitForChanges } = await newSpecPage({
+      components: [CmpA],
+    });
+
+    const cmp = doc.createElement('cmp-a');
+    doc.body.appendChild(cmp);
+
+    await waitForChanges();
+
+    // Create a container we will move the component to
+    const container = doc.createElement('div');
+    doc.body.appendChild(container);
+
+    // Move the component
+    container.appendChild(cmp);
+
+    container.remove();
+
+    expect(connected).toEqual(2);
+    expect(disconnected).toEqual(2);
+  });
 });