From 907be0ac4ff11a95e2bbc0a40e0b24885178bd85 Mon Sep 17 00:00:00 2001
From: Michael Cousins <michael@cousins.io>
Date: Sun, 9 Mar 2025 14:24:05 -0400
Subject: [PATCH] refactor(core): consolidate options validation and cleanup
 into core

---
 src/core/cleanup.js             | 32 +++++++++++++++
 src/core/index.js               | 20 ++--------
 src/core/mount-legacy.js        | 45 ++++++++++-----------
 src/core/mount-modern.svelte.js | 36 +++++++----------
 src/core/mount.js               | 36 +++++++++++++++++
 src/core/prepare.js             | 55 ++++++++++++++++++++++----
 src/core/types.d.ts             |  5 +++
 src/index.js                    |  7 ++--
 src/pure.js                     | 70 +++++++--------------------------
 9 files changed, 181 insertions(+), 125 deletions(-)
 create mode 100644 src/core/cleanup.js
 create mode 100644 src/core/mount.js

diff --git a/src/core/cleanup.js b/src/core/cleanup.js
new file mode 100644
index 0000000..2aa70da
--- /dev/null
+++ b/src/core/cleanup.js
@@ -0,0 +1,32 @@
+/** @type {Set<() => void>} */
+const cleanupTasks = new Set()
+
+/**
+ * Register later cleanup task
+ *
+ * @param {() => void} onCleanup
+ */
+const addCleanupTask = (onCleanup) => {
+  cleanupTasks.add(onCleanup)
+  return onCleanup
+}
+
+/**
+ * Remove a cleanup task without running it.
+ *
+ * @param {() => void} onCleanup
+ */
+const removeCleanupTask = (onCleanup) => {
+  cleanupTasks.delete(onCleanup)
+}
+
+/** Clean up all components and elements added to the document. */
+const cleanup = () => {
+  for (const handleCleanup of cleanupTasks.values()) {
+    handleCleanup()
+  }
+
+  cleanupTasks.clear()
+}
+
+export { addCleanupTask, cleanup, removeCleanupTask }
diff --git a/src/core/index.js b/src/core/index.js
index 1c638e8..177c3f5 100644
--- a/src/core/index.js
+++ b/src/core/index.js
@@ -5,20 +5,6 @@
  * Will switch to legacy, class-based mounting logic
  * if it looks like we're in a Svelte <= 4 environment.
  */
-import * as MountLegacy from './mount-legacy.js'
-import * as MountModern from './mount-modern.svelte.js'
-import { createValidateOptions, UnknownSvelteOptionsError } from './prepare.js'
-
-const { mount, unmount, updateProps, allowedOptions } =
-  MountModern.IS_MODERN_SVELTE ? MountModern : MountLegacy
-
-/** Validate component options. */
-const validateOptions = createValidateOptions(allowedOptions)
-
-export {
-  mount,
-  UnknownSvelteOptionsError,
-  unmount,
-  updateProps,
-  validateOptions,
-}
+export { cleanup } from './cleanup.js'
+export { mount } from './mount.js'
+export { prepare, UnknownSvelteOptionsError } from './prepare.js'
diff --git a/src/core/mount-legacy.js b/src/core/mount-legacy.js
index c9e6d1c..f2927c5 100644
--- a/src/core/mount-legacy.js
+++ b/src/core/mount-legacy.js
@@ -4,8 +4,10 @@
  * Supports Svelte <= 4.
  */
 
+import { addCleanupTask, removeCleanupTask } from './cleanup.js'
+
 /** Allowed options for the component constructor. */
-const allowedOptions = [
+const ALLOWED_OPTIONS = [
   'target',
   'accessors',
   'anchor',
@@ -15,32 +17,31 @@ const allowedOptions = [
   'context',
 ]
 
-/**
- * Mount the component into the DOM.
- *
- * The `onDestroy` callback is included for strict backwards compatibility
- * with previous versions of this library. It's mostly unnecessary logic.
- */
-const mount = (Component, options, onDestroy) => {
+/** Mount the component into the DOM. */
+const mount = (Component, options) => {
   const component = new Component(options)
 
-  if (typeof onDestroy === 'function') {
-    component.$$.on_destroy.push(() => {
-      onDestroy(component)
-    })
+  /** Remove the component from the DOM. */
+  const unmount = () => {
+    component.$destroy()
+    removeCleanupTask(unmount)
   }
 
-  return component
-}
+  /** Update the component's props. */
+  const rerender = (nextProps) => {
+    component.$set(nextProps)
+  }
 
-/** Remove the component from the DOM. */
-const unmount = (component) => {
-  component.$destroy()
-}
+  // This `$$.on_destroy` listener is included for strict backwards compatibility
+  // with previous versions of `@testing-library/svelte`.
+  // It's unnecessary and will be removed in a future major version.
+  component.$$.on_destroy.push(() => {
+    removeCleanupTask(unmount)
+  })
+
+  addCleanupTask(unmount)
 
-/** Update the component's props. */
-const updateProps = (component, nextProps) => {
-  component.$set(nextProps)
+  return { component, unmount, rerender }
 }
 
-export { allowedOptions, mount, unmount, updateProps }
+export { ALLOWED_OPTIONS, mount }
diff --git a/src/core/mount-modern.svelte.js b/src/core/mount-modern.svelte.js
index 34893f5..abd353b 100644
--- a/src/core/mount-modern.svelte.js
+++ b/src/core/mount-modern.svelte.js
@@ -5,14 +5,13 @@
  */
 import * as Svelte from 'svelte'
 
-/** Props signals for each rendered component. */
-const propsByComponent = new Map()
+import { addCleanupTask, removeCleanupTask } from './cleanup.js'
 
 /** Whether we're using Svelte >= 5. */
 const IS_MODERN_SVELTE = typeof Svelte.mount === 'function'
 
 /** Allowed options to the `mount` call. */
-const allowedOptions = [
+const ALLOWED_OPTIONS = [
   'target',
   'anchor',
   'props',
@@ -26,26 +25,21 @@ const mount = (Component, options) => {
   const props = $state(options.props ?? {})
   const component = Svelte.mount(Component, { ...options, props })
 
-  Svelte.flushSync()
-  propsByComponent.set(component, props)
+  /** Remove the component from the DOM. */
+  const unmount = () => {
+    Svelte.flushSync(() => Svelte.unmount(component))
+    removeCleanupTask(unmount)
+  }
 
-  return component
-}
+  /** Update the component's props. */
+  const rerender = (nextProps) => {
+    Svelte.flushSync(() => Object.assign(props, nextProps))
+  }
 
-/** Remove the component from the DOM. */
-const unmount = (component) => {
-  propsByComponent.delete(component)
-  Svelte.flushSync(() => Svelte.unmount(component))
-}
+  addCleanupTask(unmount)
+  Svelte.flushSync()
 
-/**
- * Update the component's props.
- *
- * Relies on the `$state` signal added in `mount`.
- */
-const updateProps = (component, nextProps) => {
-  const prevProps = propsByComponent.get(component)
-  Object.assign(prevProps, nextProps)
+  return { component, unmount, rerender }
 }
 
-export { allowedOptions, IS_MODERN_SVELTE, mount, unmount, updateProps }
+export { ALLOWED_OPTIONS, IS_MODERN_SVELTE, mount }
diff --git a/src/core/mount.js b/src/core/mount.js
new file mode 100644
index 0000000..9236ce0
--- /dev/null
+++ b/src/core/mount.js
@@ -0,0 +1,36 @@
+import { tick } from 'svelte'
+
+import * as MountLegacy from './mount-legacy.js'
+import * as MountModern from './mount-modern.svelte.js'
+
+const mountComponent = MountModern.IS_MODERN_SVELTE
+  ? MountModern.mount
+  : MountLegacy.mount
+
+/**
+ * Render a Svelte component into the document.
+ *
+ * @template {import('./types.js').Component} C
+ * @param {import('./types.js').ComponentType<C>} Component
+ * @param {import('./types.js').MountOptions<C>} options
+ * @returns {{
+ *   component: C
+ *   unmount: () => void
+ *   rerender: (props: Partial<import('./types.js').Props<C>>) => Promise<void>
+ * }}
+ */
+const mount = (Component, options = {}) => {
+  const { component, unmount, rerender } = mountComponent(Component, options)
+
+  return {
+    component,
+    unmount,
+    rerender: async (props) => {
+      rerender(props)
+      // Await the next tick for Svelte 4, which cannot flush changes synchronously
+      await tick()
+    },
+  }
+}
+
+export { mount }
diff --git a/src/core/prepare.js b/src/core/prepare.js
index c0d794b..0a91270 100644
--- a/src/core/prepare.js
+++ b/src/core/prepare.js
@@ -1,9 +1,18 @@
+import { addCleanupTask } from './cleanup.js'
+import * as MountLegacy from './mount-legacy.js'
+import * as MountModern from './mount-modern.svelte.js'
+
+const ALLOWED_OPTIONS = MountModern.IS_MODERN_SVELTE
+  ? MountModern.ALLOWED_OPTIONS
+  : MountLegacy.ALLOWED_OPTIONS
+
+/** An error thrown for incorrect options and clashes between props and Svelte options. */
 class UnknownSvelteOptionsError extends TypeError {
-  constructor(unknownOptions, allowedOptions) {
+  constructor(unknownOptions) {
     super(`Unknown options.
 
     Unknown: [ ${unknownOptions.join(', ')} ]
-    Allowed: [ ${allowedOptions.join(', ')} ]
+    Allowed: [ ${ALLOWED_OPTIONS.join(', ')} ]
 
     To pass both Svelte options and props to a component,
     or to use props that share a name with a Svelte option,
@@ -15,9 +24,41 @@ class UnknownSvelteOptionsError extends TypeError {
   }
 }
 
-const createValidateOptions = (allowedOptions) => (options) => {
+/**
+ * Prepare DOM elements for rendering.
+ *
+ * @template {import('./types.js').Component} C
+ * @param {import('./types.js').PropsOrMountOptions<C>} propsOrOptions
+ * @param {{ baseElement?: HTMLElement }} renderOptions
+ * @returns {{
+ *   baseElement: HTMLElement
+ *   target: HTMLElement
+ *   mountOptions: import('./types.js').MountOptions<C>
+ * }}
+ */
+const prepare = (propsOrOptions = {}, renderOptions = {}) => {
+  const mountOptions = validateMountOptions(propsOrOptions)
+
+  const baseElement =
+    renderOptions.baseElement ?? mountOptions.target ?? document.body
+
+  const target =
+    mountOptions.target ??
+    baseElement.appendChild(document.createElement('div'))
+
+  addCleanupTask(() => {
+    if (target.parentNode === document.body) {
+      document.body.removeChild(target)
+    }
+  })
+
+  return { baseElement, target, mountOptions: { ...mountOptions, target } }
+}
+
+/** Prevent incorrect options and clashes between props and Svelte options. */
+const validateMountOptions = (options) => {
   const isProps = !Object.keys(options).some((option) =>
-    allowedOptions.includes(option)
+    ALLOWED_OPTIONS.includes(option)
   )
 
   if (isProps) {
@@ -26,14 +67,14 @@ const createValidateOptions = (allowedOptions) => (options) => {
 
   // Check if any props and Svelte options were accidentally mixed.
   const unknownOptions = Object.keys(options).filter(
-    (option) => !allowedOptions.includes(option)
+    (option) => !ALLOWED_OPTIONS.includes(option)
   )
 
   if (unknownOptions.length > 0) {
-    throw new UnknownSvelteOptionsError(unknownOptions, allowedOptions)
+    throw new UnknownSvelteOptionsError(unknownOptions)
   }
 
   return options
 }
 
-export { createValidateOptions, UnknownSvelteOptionsError }
+export { prepare, UnknownSvelteOptionsError }
diff --git a/src/core/types.d.ts b/src/core/types.d.ts
index 8300243..0023158 100644
--- a/src/core/types.d.ts
+++ b/src/core/types.d.ts
@@ -59,3 +59,8 @@ export type Exports<C> = IS_MODERN_SVELTE extends true
 export type MountOptions<C extends Component> = IS_MODERN_SVELTE extends true
   ? Parameters<typeof mount<Props<C>, Exports<C>>>[1]
   : LegacyConstructorOptions<Props<C>>
+
+/** Component props or partial mount options. */
+export type PropsOrMountOptions<C extends Component> =
+  | Props<C>
+  | Partial<MountOptions<C>>
diff --git a/src/index.js b/src/index.js
index 2704824..fbfecb9 100644
--- a/src/index.js
+++ b/src/index.js
@@ -1,5 +1,6 @@
 /* eslint-disable import/export */
-import { act, cleanup } from './pure.js'
+import { cleanup } from './core/index.js'
+import { act } from './pure.js'
 
 // If we're running in a test runner that supports afterEach
 // then we'll automatically run cleanup afterEach test
@@ -16,7 +17,7 @@ if (typeof afterEach === 'function' && !process.env.STL_SKIP_AUTO_CLEANUP) {
 export * from '@testing-library/dom'
 
 // export svelte-specific functions and custom `fireEvent`
-export { UnknownSvelteOptionsError } from './core/index.js'
-export * from './pure.js'
 // `fireEvent` must be named to take priority over wildcard from @testing-library/dom
+export { cleanup, UnknownSvelteOptionsError } from './core/index.js'
 export { fireEvent } from './pure.js'
+export * from './pure.js'
diff --git a/src/pure.js b/src/pure.js
index 35a10fa..78f1e79 100644
--- a/src/pure.js
+++ b/src/pure.js
@@ -5,16 +5,13 @@ import {
 } from '@testing-library/dom'
 import { tick } from 'svelte'
 
-import { mount, unmount, updateProps, validateOptions } from './core/index.js'
-
-const targetCache = new Set()
-const componentCache = new Set()
+import { mount, prepare } from './core/index.js'
 
 /**
  * Customize how Svelte renders the component.
  *
  * @template {import('./core/types.js').Component} C
- * @typedef {import('./core/types.js').Props<C> | Partial<import('./core/types.js').MountOptions<C>>} SvelteComponentOptions
+ * @typedef {import('./core/types.js').PropsOrMountOptions<C>} SvelteComponentOptions
  */
 
 /**
@@ -52,38 +49,28 @@ const componentCache = new Set()
  * @template {import('@testing-library/dom').Queries} [Q=typeof import('@testing-library/dom').queries]
  *
  * @param {import('./core/types.js').ComponentType<C>} Component - The component to render.
- * @param {SvelteComponentOptions<C>} options - Customize how Svelte renders the component.
+ * @param {SvelteComponentOptions<C>} propsOrOptions - Customize how Svelte renders the component.
  * @param {RenderOptions<Q>} renderOptions - Customize how Testing Library sets up the document and binds queries.
  * @returns {RenderResult<C, Q>} The rendered component and bound testing functions.
  */
-const render = (Component, options = {}, renderOptions = {}) => {
-  options = validateOptions(options)
-
-  const baseElement =
-    renderOptions.baseElement ?? options.target ?? document.body
-
-  const queries = getQueriesForElement(baseElement, renderOptions.queries)
-
-  const target =
-    options.target ?? baseElement.appendChild(document.createElement('div'))
-
-  targetCache.add(target)
+const render = (Component, propsOrOptions = {}, renderOptions = {}) => {
+  const { baseElement, target, mountOptions } = prepare(
+    propsOrOptions,
+    renderOptions
+  )
 
-  const component = mount(
+  const { component, unmount, rerender } = mount(
     Component.default ?? Component,
-    { ...options, target },
-    cleanupComponent
+    mountOptions
   )
 
-  componentCache.add(component)
+  const queries = getQueriesForElement(baseElement, renderOptions.queries)
 
   return {
     baseElement,
     component,
     container: target,
-    debug: (el = baseElement) => {
-      console.log(prettyDOM(el))
-    },
+    debug: (el = baseElement) => console.log(prettyDOM(el)),
     rerender: async (props) => {
       if (props.props) {
         console.warn(
@@ -92,40 +79,13 @@ const render = (Component, options = {}, renderOptions = {}) => {
         props = props.props
       }
 
-      updateProps(component, props)
-      await tick()
-    },
-    unmount: () => {
-      cleanupComponent(component)
+      await rerender(props)
     },
+    unmount,
     ...queries,
   }
 }
 
-/** Remove a component from the component cache. */
-const cleanupComponent = (component) => {
-  const inCache = componentCache.delete(component)
-
-  if (inCache) {
-    unmount(component)
-  }
-}
-
-/** Remove a target element from the target cache. */
-const cleanupTarget = (target) => {
-  const inCache = targetCache.delete(target)
-
-  if (inCache && target.parentNode === document.body) {
-    document.body.removeChild(target)
-  }
-}
-
-/** Unmount all components and remove elements added to `<body>`. */
-const cleanup = () => {
-  componentCache.forEach(cleanupComponent)
-  targetCache.forEach(cleanupTarget)
-}
-
 /**
  * Call a function and wait for Svelte to flush pending changes.
  *
@@ -171,4 +131,4 @@ Object.keys(baseFireEvent).forEach((key) => {
   }
 })
 
-export { act, cleanup, fireEvent, render }
+export { act, fireEvent, render }