hotwired · seanpdoyle · Sep 13, 2022
diff --git a/src/core/drive/page_renderer.ts b/src/core/drive/page_renderer.ts
@@ -1,19 +1,22 @@
 import { Renderer } from "../renderer"
 import { PageSnapshot } from "./page_snapshot"
 import { ReloadReason } from "../native/browser_adapter"
-import { activateScriptElement, waitForLoad } from "../../util"
+import { activateScriptElement, waitForLoad, getBodyElementId } from "../../util"
 
 export class PageRenderer extends Renderer<HTMLBodyElement, PageSnapshot> {
-  static renderElement(currentElement: HTMLBodyElement, newElement: HTMLBodyElement) {
-    if (document.body && newElement instanceof HTMLBodyElement) {
-      document.body.replaceWith(newElement)
+  static renderElement(currentBody: HTMLBodyElement, newBody: HTMLBodyElement) {
+    if (document.body && newBody instanceof HTMLBodyElement) {
+      const currentElement = PageRenderer.getRenderedElement(currentBody) || currentBody
+      const newElement = PageRenderer.getRenderedElement(newBody) || newBody
+
+      currentElement.replaceWith(newElement)
     } else {
-      document.documentElement.appendChild(newElement)
+      document.documentElement.appendChild(newBody)
     }
   }
 
   get shouldRender() {
-    return this.newSnapshot.isVisitable && this.trackedElementsAreIdentical
+    return this.newSnapshot.isVisitable && this.trackedElementsAreIdentical && this.renderedElementMatches
   }
 
   get reloadReason(): ReloadReason {
@@ -28,6 +31,12 @@ export class PageRenderer extends Renderer<HTMLBodyElement, PageSnapshot> {
         reason: "tracked_element_mismatch",
       }
     }
+
+    if (!this.renderedElementMatches) {
+      return {
+        reason: "rendered_element_mismatch",
+      }
+    }
   }
 
   async prepareToRender() {
@@ -78,6 +87,24 @@ export class PageRenderer extends Renderer<HTMLBodyElement, PageSnapshot> {
     return this.currentHeadSnapshot.trackedElementSignature == this.newHeadSnapshot.trackedElementSignature
   }
 
+  get renderedElementMatches() {
+    return PageRenderer.getRenderedElement(this.newElement) !== null
+  }
+
+  static get bodySelector() {
+    const bodyId = getBodyElementId()
+
+    return bodyId ? `#${bodyId}` : "body"
+  }
+
+  static getRenderedElement(element: HTMLElement): HTMLElement | null {
+    if (element.matches(this.bodySelector)) {
+      return element
+    } else {
+      return element.querySelector(this.bodySelector)
+    }
+  }
+
   async copyNewHeadStylesheetElements() {
     const loadingElements = []
 

diff --git a/src/core/drive/page_view.ts b/src/core/drive/page_view.ts
@@ -1,4 +1,3 @@
-import { nextEventLoopTick } from "../../util"
 import { View, ViewDelegate, ViewRenderOptions } from "../view"
 import { ErrorRenderer } from "./error_renderer"
 import { PageRenderer } from "./page_renderer"
@@ -39,11 +38,10 @@ export class PageView extends View<HTMLBodyElement, PageSnapshot, PageViewRender
     this.snapshotCache.clear()
   }
 
-  async cacheSnapshot() {
+  cacheSnapshot() {
     if (this.shouldCacheSnapshot) {
       this.delegate.viewWillCacheSnapshot()
       const { snapshot, lastRenderedLocation: location } = this
-      await nextEventLoopTick()
       const cachedSnapshot = snapshot.clone()
       this.snapshotCache.put(location, cachedSnapshot)
       return cachedSnapshot

diff --git a/src/core/drive/visit.ts b/src/core/drive/visit.ts
@@ -465,7 +465,8 @@ export class Visit implements FetchRequestDelegate {
 
   cacheSnapshot() {
     if (!this.snapshotCached) {
-      this.view.cacheSnapshot().then((snapshot) => snapshot && this.visitCachedSnapshot(snapshot))
+      const snapshot = this.view.cacheSnapshot()
+      if (snapshot) this.visitCachedSnapshot(snapshot)
       this.snapshotCached = true
     }
   }

diff --git a/src/tests/fixtures/drive_custom_body.html b/src/tests/fixtures/drive_custom_body.html
@@ -0,0 +1,21 @@
+<!DOCTYPE html>
+<html>
+  <head>
+    <meta charset="utf-8">
+    <meta name="turbo-body" content="app">
+    <title>Drive (with custom body)</title>
+    <script src="/dist/turbo.es2017-umd.js" data-turbo-track="reload"></script>
+    <script src="/src/tests/fixtures/test.js"></script>
+  </head>
+  <body>
+    <h1>Drive (with custom body)</h1>
+    <div id="app">
+      <div>
+        <a id="drive" href="/src/tests/fixtures/drive_custom_body_2.html">Drive enabled link</a>
+        <a id="mismatch" href="/src/tests/fixtures/drive_custom_body_3.html">Drive enabled link to page with mismatched custom body</a>
+      </div>
+
+      <p id="different-content">Drive 1</p>
+    </div>
+  </body>
+</html>
diff --git a/src/tests/fixtures/drive_custom_body_2.html b/src/tests/fixtures/drive_custom_body_2.html
@@ -0,0 +1,21 @@
+<!DOCTYPE html>
+<html>
+  <head>
+    <meta charset="utf-8">
+    <meta name="turbo-body" content="app">
+    <title>Drive (with custom body)</title>
+    <script src="/dist/turbo.es2017-umd.js" data-turbo-track="reload"></script>
+    <script src="/src/tests/fixtures/test.js"></script>
+  </head>
+  <body>
+    <h1>Drive (with custom body 2)</h1>
+
+    <div id="app">
+      <div>
+        <a id="drive" href="/src/tests/fixtures/drive_custom_body.html">Drive enabled link</a>
+      </div>
+
+      <p id="different-content">Drive 2</p>
+    </div>
+  </body>
+</html>
diff --git a/src/tests/fixtures/drive_custom_body_3.html b/src/tests/fixtures/drive_custom_body_3.html
@@ -0,0 +1,15 @@
+<!DOCTYPE html>
+<html>
+  <head>
+    <meta charset="utf-8">
+    <meta name="turbo-body" content="main">
+    <title>Drive (with custom body)</title>
+    <script src="/dist/turbo.es2017-umd.js" data-turbo-track="reload"></script>
+    <script src="/src/tests/fixtures/test.js"></script>
+  </head>
+  <body>
+    <main id="main">
+      <h1>Drive (with custom body 3)</h1>
+    </main>
+  </body>
+</html>
diff --git a/src/tests/functional/drive_custom_body_tests.ts b/src/tests/functional/drive_custom_body_tests.ts
@@ -0,0 +1,40 @@
+import { test } from "@playwright/test"
+import { assert } from "chai"
+import { nextEventNamed, pathname, getFromLocalStorage, setLocalStorageFromEvent } from "../helpers/page"
+
+test.beforeEach(async ({ page }) => {
+  await page.goto("/src/tests/fixtures/drive_custom_body.html")
+})
+
+test("test drive with a custom body element", async ({ page }) => {
+  await page.click("#drive")
+  await nextEventNamed(page, "turbo:load")
+
+  assert.equal(pathname(page.url()), "/src/tests/fixtures/drive_custom_body_2.html")
+  assert.equal(await page.textContent("h1"), "Drive (with custom body)")
+  assert.equal(await page.textContent("#different-content"), "Drive 2")
+
+  await page.goBack()
+  await nextEventNamed(page, "turbo:load")
+
+  assert.equal(pathname(page.url()), "/src/tests/fixtures/drive_custom_body.html")
+  assert.equal(await page.textContent("h1"), "Drive (with custom body)")
+  assert.equal(await page.textContent("#different-content"), "Drive 1")
+
+  await page.goForward()
+  await nextEventNamed(page, "turbo:load")
+
+  assert.equal(pathname(page.url()), "/src/tests/fixtures/drive_custom_body_2.html")
+  assert.equal(await page.textContent("h1"), "Drive (with custom body)")
+  assert.equal(await page.textContent("#different-content"), "Drive 2")
+})
+
+test("test drive with mismatched custom body elements", async ({ page }) => {
+  await setLocalStorageFromEvent(page, "turbo:reload", "reloaded", "true")
+  await page.click("#mismatch")
+  await page.waitForEvent("load")
+
+  assert.equal(await page.textContent("h1"), "Drive (with custom body 3)")
+  assert.equal(await getFromLocalStorage(page, "reloaded"), "true", "dispatches turbo:reload event")
+  assert.equal(pathname(page.url()), "/src/tests/fixtures/drive_custom_body_3.html")
+})
diff --git a/src/util.ts b/src/util.ts
@@ -160,6 +160,10 @@ export function getVisitAction(...elements: (Element | undefined)[]): Action | n
   return isAction(action) ? action : null
 }
 
+export function getBodyElementId(): string | null {
+  return getMetaContent("turbo-body")
+}
+
 export function getMetaElement(name: string): HTMLMetaElement | null {
   return document.querySelector(`meta[name="${name}"]`)
 }