Expose custom return type for our mount function

src/mount.ts

@@ -1,6 +1,12 @@
 import * as enzyme from 'enzyme';
 import type { ReactWrapper } from 'enzyme';
-import type { VNode } from 'preact';
+import {
+  Component,
+  ComponentClass,
+  FunctionComponent,
+  JSX,
+  VNode,
+} from 'preact';
 
 let containers: HTMLElement[] = [];
 let wrappers: ReactWrapper[] = [];
@@ -20,16 +26,80 @@ export type MountOptions = {
   prepareContainer?: (container: HTMLElement) => void;
 };
 
+export type EnzymePropSelector = {
+  [key: string]: any;
+};
+
+export type ComponentType<Props> =
+  | ComponentClass<Props>
+  | FunctionComponent<Props>;
+
+export type EnzymeSelector = string | ComponentType<any> | EnzymePropSelector;
+
+/**
+ * Inspired by DefinitelyTyped's ReactWrapper.
+ * @see https://github.com/DefinitelyTyped/DefinitelyTyped/blob/cebb88fecfa52f854826e216a537867c11b0150e/types/enzyme/index.d.ts
+ */
+export type ComponentWrapper<P = {}, S = {}, C = Component> = {
+  at(index: number): ComponentWrapper;
+
+  childAt(index: number): ComponentWrapper<any, any>;
+  childAt<P2, S2>(index: number): ComponentWrapper<P2, S2>;
+
+  closest<P2>(
+    statelessComponent: FunctionComponent<P2>,
+  ): ComponentWrapper<P2, never>;
+  closest<P2>(component: ComponentType<P2>): ComponentWrapper<P2, any>;
+  closest(props: EnzymePropSelector): ComponentWrapper<any, any>;
+  closest(selector: string): ComponentWrapper<JSX.HTMLAttributes, any>;
+
+  exists(selector?: EnzymeSelector): boolean;
+
+  find<P2>(
+    statelessComponent: FunctionComponent<P2>,
+  ): ComponentWrapper<P2, never>;
+  find<P2>(component: ComponentType<P2>): ComponentWrapper<P2, any>;
+  find<C2 extends Component>(
+    componentClass: ComponentClass<C2['props']>,
+  ): ComponentWrapper<C2['props'], C2['state'], C2>;
+  find(props: EnzymePropSelector): ComponentWrapper<any, any>;
+  find(selector: string): ComponentWrapper<JSX.HTMLAttributes, any>;
+
+  first(): ComponentWrapper<P, S, C>;
+  forEach(
+    fn: (wrapper: ComponentWrapper<P, S, C>, index: number) => any,
+  ): ComponentWrapper<P, S, C>;
+  getDOMNode(): HTMLElement;
+  hasClass(className: string | RegExp): boolean;
+  instance(): C;
+  last(): ComponentWrapper<P, S, C>;
+  length: number;
+  map<V>(fn: (wrapper: ComponentWrapper<P, S, C>, index: number) => V): V[];
+
+  prop<K extends keyof P>(key: K): P[K];
+  prop<T>(key: string): T;
+
+  props(): P;
+  setProps<K extends keyof P>(
+    props: Pick<P, K>,
+    callback?: () => void,
+  ): ComponentWrapper<P, S, C>;
+  simulate(event: string, ...args: any[]): ComponentWrapper<P, S, C>;
+  text(): string;
+  unmount(): void;
+  update(): ComponentWrapper<P, S, C>;
+};
+
 /**
  * Render a Preact component using Enzyme and return a wrapper.
  *
  * The component can be unmounted by calling `wrapper.unmount()` or by calling
  * {@link unmountAll} at the end of the test.
  */
-export function mount(
+export function mount<P = {}, S = {}, C = Component>(
   jsx: VNode,
   { connected = false, prepareContainer }: MountOptions = {},
-) {
+): ComponentWrapper<P, S, C> {
   let wrapper;
   if (connected) {
     const container = document.createElement('div');
@@ -46,7 +116,7 @@ export function mount(
 
   wrappers.push(wrapper);
 
-  return wrapper;
+  return wrapper as ComponentWrapper<P, S, C>;
 }
 
 /**