From 9f9f234ed2090bbc88aa58c54d12a1c534d8a8bf Mon Sep 17 00:00:00 2001
From: Chris Villa <chrisvilla@me.com>
Date: Tue, 8 Sep 2020 12:39:25 +0100
Subject: [PATCH] feat: support custom rehydratable query selectors

---
 src/ILoadedOptions.ts                     |  7 ++
 src/IOptions.ts                           |  1 +
 src/__tests__/__snapshots__/tests.ts.snap |  2 +
 src/__tests__/tests.ts                    | 20 ++++++
 src/rehydrator.ts                         | 78 ++++++++++++++++-------
 5 files changed, 86 insertions(+), 22 deletions(-)
 create mode 100644 src/ILoadedOptions.ts

diff --git a/src/ILoadedOptions.ts b/src/ILoadedOptions.ts
new file mode 100644
index 0000000..dd10fd2
--- /dev/null
+++ b/src/ILoadedOptions.ts
@@ -0,0 +1,7 @@
+interface IOptions {
+  allSelectors: { [key: string]: string };
+  compoundSelector: string;
+  extra: object;
+}
+
+export default IOptions;
diff --git a/src/IOptions.ts b/src/IOptions.ts
index 888b0d6..404ac44 100644
--- a/src/IOptions.ts
+++ b/src/IOptions.ts
@@ -1,5 +1,6 @@
 interface IOptions {
   extra: object;
+  getQuerySelector?: (key: string) => string;
 }
 
 export default IOptions;
diff --git a/src/__tests__/__snapshots__/tests.ts.snap b/src/__tests__/__snapshots__/tests.ts.snap
index cd4d07f..60d3b23 100644
--- a/src/__tests__/__snapshots__/tests.ts.snap
+++ b/src/__tests__/__snapshots__/tests.ts.snap
@@ -2,6 +2,8 @@
 
 exports[`reactFromHtml E2E tests Should rehydrate a basic component 1`] = `"<div class=\\"rehydration-root\\"><span>rehydrated component</span></div>"`;
 
+exports[`reactFromHtml E2E tests Should rehydrate components with custom query selectors 1`] = `"<div class=\\"rehydration-root\\"><span>rehydrated component</span></div>"`;
+
 exports[`reactFromHtml E2E tests Should work for nested rehydratables 1`] = `
 "
       <div class=\\"rehydration-root\\"><span>
diff --git a/src/__tests__/tests.ts b/src/__tests__/tests.ts
index a91709b..5d59369 100644
--- a/src/__tests__/tests.ts
+++ b/src/__tests__/tests.ts
@@ -51,4 +51,24 @@ describe("reactFromHtml E2E tests", async () => {
     expect(documentElement.innerHTML).toMatchSnapshot();
     expect(mockCall).toBeCalledTimes(2);
   });
+
+  it("Should rehydrate components with custom query selectors", async () => {
+    const componentName: string = "myComponent";
+
+    const rehydrator = async () => {
+      return React.createElement("span", {}, "rehydrated component");
+    };
+
+    const rehydrators = { [componentName]: rehydrator };
+    const documentElement = document.createElement("div");
+
+    documentElement.innerHTML = `<div class="test-${componentName}"></div>`;
+
+    await reactFromHtml(documentElement, rehydrators, {
+      extra: {},
+      getQuerySelector: key => `.test-${key}`,
+    });
+
+    expect(documentElement.innerHTML).toMatchSnapshot();
+  });
 });
diff --git a/src/rehydrator.ts b/src/rehydrator.ts
index 72426e9..2dd0465 100644
--- a/src/rehydrator.ts
+++ b/src/rehydrator.ts
@@ -1,15 +1,24 @@
 import * as ReactDOM from "react-dom";
 
 import domElementToReact from "./dom-element-to-react";
+import ILoadedOptions from "./ILoadedOptions";
 import IOptions from "./IOptions";
 import IRehydrator from "./IRehydrator";
 
 const rehydratableToReactElement = async (
   el: Element,
   rehydrators: IRehydrator,
-  options: IOptions
+  options: ILoadedOptions
 ): Promise<React.ReactElement<any>> => {
-  const rehydratorName = el.getAttribute("data-rehydratable");
+  const rehydratorSelector = Object.keys(options.allSelectors).find(selector =>
+    el.matches(selector)
+  );
+
+  if (!rehydratorSelector) {
+    throw new Error("No rehydrator selector matched the element.");
+  }
+
+  const rehydratorName = options.allSelectors[rehydratorSelector];
 
   if (!rehydratorName) {
     throw new Error("Rehydrator name is missing from element.");
@@ -33,13 +42,13 @@ const rehydratableToReactElement = async (
 
 const createCustomHandler = (
   rehydrators: IRehydrator,
-  options: IOptions
+  options: ILoadedOptions
 ) => async (node: Node) => {
   // This function will run on _every_ node that domElementToReact encounters.
   // Make sure to keep the conditional highly performant.
   if (
     node.nodeType === Node.ELEMENT_NODE &&
-    (node as Element).hasAttribute("data-rehydratable")
+    (node as Element).matches(options.compoundSelector)
   ) {
     return rehydratableToReactElement(node as Element, rehydrators, options);
   }
@@ -63,7 +72,7 @@ const createReactRoot = (el: Node) => {
 const rehydrateChildren = async (
   el: Node,
   rehydrators: IRehydrator,
-  options: IOptions
+  options: ILoadedOptions
 ) => {
   const container = createReactRoot(el);
 
@@ -93,30 +102,55 @@ const render = ({
   ReactDOM.render(rehydrated as React.ReactElement<any>, root);
 };
 
-const createQuerySelector = (rehydratableIds: string[]) =>
-  rehydratableIds.reduce(
-    (acc: string, rehydratableId: string) =>
-      `${acc ? `${acc}, ` : ""}[data-rehydratable*="${rehydratableId}"]`,
+const defaultGetQuerySelector = (key: string) =>
+  `[data-rehydratable*="${key}"]`;
+
+const createQuerySelectors = (
+  rehydratableIds: string[],
+  getQuerySelector: ((key: string) => string) = defaultGetQuerySelector
+) => {
+  const allSelectors: { [key: string]: string } = rehydratableIds.reduce(
+    (acc, key) => ({ ...acc, [getQuerySelector(key)]: key }),
+    {}
+  );
+
+  const compoundSelector = Object.keys(allSelectors).reduce(
+    (acc: string, selector: string) => `${acc ? `${acc}, ` : ""}${selector}`,
     ""
   );
 
+  return {
+    allSelectors,
+    compoundSelector,
+  };
+};
+
 const rehydrate = async (
   container: Element,
   rehydrators: IRehydrator,
   options: IOptions
 ) => {
-  const selector = createQuerySelector(Object.keys(rehydrators));
-
-  const roots = Array.from(
-    // TODO: allow setting a container identifier so multiple rehydration instances can exist
-    container.querySelectorAll(selector)
-  ).reduce((acc: Element[], root: Element) => {
-    // filter roots that are contained within other roots
-    if (!acc.some(r => r.contains(root))) {
-      acc.push(root);
-    }
-    return acc;
-  }, []);
+  const { allSelectors, compoundSelector } = createQuerySelectors(
+    Object.keys(rehydrators),
+    options.getQuerySelector
+  );
+
+  const loadedOptions: ILoadedOptions = {
+    allSelectors,
+    compoundSelector,
+    extra: options.extra,
+  };
+
+  const roots = Array.from(container.querySelectorAll(compoundSelector)).reduce(
+    (acc: Element[], root: Element) => {
+      // filter roots that are contained within other roots
+      if (!acc.some(r => r.contains(root))) {
+        acc.push(root);
+      }
+      return acc;
+    },
+    []
+  );
 
   // TODO: solve race condition when a second rehydrate runs
 
@@ -130,7 +164,7 @@ const rehydrate = async (
           const {
             container: rootContainer,
             rehydrated,
-          } = await rehydrateChildren(root, rehydrators, options);
+          } = await rehydrateChildren(root, rehydrators, loadedOptions);
 
           return { root: rootContainer, rehydrated };
         } catch (e) {