fix(hydrate): ensure beforeHydrateFn and afterHydrateFn always return a function (#5890)

christian-bromann · web-flow · commit a7c212c2a9de · 2024-07-08T18:22:36.000Z
diff --git a/src/hydrate/runner/render.ts b/src/hydrate/runner/render.ts
@@ -130,7 +130,7 @@ async function render(win: MockWindow, opts: HydrateFactoryOptions, results: Hyd
   }
 
   initializeWindow(win, win.document, opts, results);
-  const beforeHydrateFn = typeof opts.beforeHydrate === 'function' ? opts.beforeHydrate(win.document) : NOOP;
+  const beforeHydrateFn = typeof opts.beforeHydrate === 'function' ? opts.beforeHydrate : NOOP;
   try {
     await Promise.resolve(beforeHydrateFn(win.document));
     return new Promise<HydrateResults>((resolve) => hydrateFactory(win, opts, results, afterHydrate, resolve));
@@ -162,7 +162,7 @@ async function afterHydrate(
   results: HydrateResults,
   resolve: (results: HydrateResults) => void,
 ) {
-  const afterHydrateFn = typeof opts.afterHydrate === 'function' ? opts.afterHydrate(win.document) : NOOP;
+  const afterHydrateFn = typeof opts.afterHydrate === 'function' ? opts.afterHydrate : NOOP;
   try {
     await Promise.resolve(afterHydrateFn(win.document));
     return resolve(finalizeHydrate(win, win.document, opts, results));
diff --git a/test/end-to-end/src/declarative-shadow-dom/test.e2e.ts b/test/end-to-end/src/declarative-shadow-dom/test.e2e.ts
@@ -253,4 +253,18 @@ describe('renderToString', () => {
       `<car-detail class=\"sc-car-list\" custom-hydrate-flag=\"\" c-id=\"2.4.2.0\" s-id=\"4\"><!--r.4--><section class=\"sc-car-list\" c-id=\"4.0.0.0\"><!--t.4.1.1.0-->2023 VW Beetle</section></car-detail>`,
     );
   });
+
+  it('calls beforeHydrate and afterHydrate function hooks', async () => {
+    const beforeHydrate = jest.fn((doc) => (doc.querySelector('div').textContent = 'Hello Universe'));
+    const afterHydrate = jest.fn();
+
+    const { html } = await renderToString('<div>Hello World</div>', {
+      beforeHydrate,
+      afterHydrate,
+    });
+
+    expect(beforeHydrate).toHaveBeenCalledTimes(1);
+    expect(afterHydrate).toHaveBeenCalledTimes(1);
+    expect(html).toContain('<body><div>Hello Universe</div></body>');
+  });
 });
