Automatically wrap the relevant methods in act

If @testing-library/react is already installed,
we use their implementation of `act` - it's complete
and provides useful warnings.
If it's not installed, then we simply assume that the user
is not testing a React application, and use a noop instead

README.md

@@ -13,7 +13,7 @@ In particular:
 
 - [testing a redux saga that manages a websocket connection](https://github.com/romgain/jest-websocket-mock/blob/master/examples/redux-saga/src/__tests__/saga.test.js)
 - [testing a component using the saga above](https://github.com/romgain/jest-websocket-mock/blob/master/examples/redux-saga/src/__tests__/App.test.js)
-- [testing a component that manages a websocket connection using react hooks](https://github.com/romgain/jest-websocket-mock/blob/master/examples/hooks/src/__tests__/App.test.js)
+- [testing a component that manages a websocket connection using react hooks](https://github.com/romgain/jest-websocket-mock/blob/master/examples/hooks/src/App.test.tsx)
 
 ## Install
 
@@ -239,6 +239,15 @@ afterEach(() => {
 });
 ```
 
+## Testing React applications
+
+When testing React applications, `jest-websocket-mock` will look for
+`@testing-library/react`'s implementation of [`act`](https://reactjs.org/docs/test-utils.html#act).
+If it is available, it will wrap all the necessary calls in `act`, so you don't have to.
+
+If `@testing-library/react` is not available, we will assume that you're not testing a React application,
+and you might need to call `act` manually.
+
 ## Using `jest-websocket-mock` to interact with a non-global WebSocket object
 
 `jest-websocket-mock` uses [Mock Socket](https://github.com/thoov/mock-socket)

examples/hooks/src/App.test.tsx

@@ -1,5 +1,5 @@
 import React from "react";
-import { act, render, screen, fireEvent } from "@testing-library/react";
+import { render, screen, fireEvent } from "@testing-library/react";
 import userEvent from "@testing-library/user-event";
 import WS from "jest-websocket-mock";
 import App from "./App";
@@ -9,30 +9,24 @@ beforeEach(() => {
   ws = new WS("ws://localhost:8080");
 });
 afterEach(() => {
-  act(() => {
-    WS.clean();
-  });
+  WS.clean();
 });
 
 describe("The App component", () => {
   it("renders a dot indicating the connection status", async () => {
     render(<App />);
     expect(screen.getByTitle("disconnected")).toBeInTheDocument();
-    await act(async () => {
-      await ws.connected;
-    });
+
+    await ws.connected;
     expect(screen.getByTitle("connected")).toBeInTheDocument();
-    act(() => {
-      ws.close();
-    });
+
+    ws.close();
     expect(screen.getByTitle("disconnected")).toBeInTheDocument();
   });
 
   it("sends and receives messages", async () => {
     render(<App />);
-    await act(async () => {
-      await ws.connected;
-    });
+    await ws.connected;
 
     const input = screen.getByPlaceholderText("type your message here...");
     userEvent.type(input, "Hello there");
@@ -41,9 +35,7 @@ describe("The App component", () => {
     await expect(ws).toReceiveMessage("Hello there");
     expect(screen.getByText("(sent) Hello there")).toBeInTheDocument();
 
-    act(() => {
-      ws.send("[echo] Hello there");
-    });
+    ws.send("[echo] Hello there");
     expect(
       screen.getByText("(received) [echo] Hello there")
     ).toBeInTheDocument();

rollup.config.js

@@ -5,7 +5,7 @@ import pkg from "./package.json";
 export default [
   {
     input: "src/index.ts",
-    external: ["mock-socket", "jest-diff"],
+    external: ["mock-socket", "jest-diff", "@testing-library/react"],
     plugins: [
       resolve({ extensions: [".js", ".ts"] }),
       babel({ extensions: [".js", ".ts"] }),

src/act-compat.ts

@@ -0,0 +1,23 @@
+/**
+ * A simple compatibility method for react's "act".
+ * If @testing-library/react is already installed, we just use
+ * their implementation - it's complete and has useful warnings.
+ * If @testing-library/react is *not* installed, then we just assume
+ * that the user is not testing a react application, and use a noop instead.
+ */
+
+type Callback = () => Promise<void | undefined> | void | undefined;
+type AsyncAct = (callback: Callback) => Promise<undefined>;
+type SyncAct = (callback: Callback) => void;
+
+let act: AsyncAct | SyncAct;
+
+try {
+  act = require("@testing-library/react").act;
+} catch (_) {
+  act = (callback: Function) => {
+    callback();
+  };
+}
+
+export default act;

src/websocket.ts

@@ -1,5 +1,6 @@
 import { Server, WebSocket, CloseOptions } from "mock-socket";
 import Queue from "./queue";
+import act from "./act-compat";
 
 const identity = (x: string) => x;
 interface WSOptions {
@@ -17,15 +18,16 @@ interface MockWebSocket extends Omit<WebSocket, "close"> {
 
 export default class WS {
   server: Server;
-  connected: Promise<WebSocket>;
-  closed: Promise<{}>;
   serializer: (deserializedMessage: DeserializedMessage) => string;
   deserializer: (message: string) => DeserializedMessage;
 
   static instances: Array<WS> = [];
   messages: Array<DeserializedMessage> = [];
   messagesToConsume = new Queue();
 
+  private _isConnected: Promise<WebSocket>;
+  private _isClosed: Promise<{}>;
+
   static clean() {
     WS.instances.forEach((instance) => {
       instance.close();
@@ -42,8 +44,8 @@ export default class WS {
 
     let connectionResolver: (socket: WebSocket) => void,
       closedResolver!: () => void;
-    this.connected = new Promise((done) => (connectionResolver = done));
-    this.closed = new Promise((done) => (closedResolver = done));
+    this._isConnected = new Promise((done) => (connectionResolver = done));
+    this._isClosed = new Promise((done) => (closedResolver = done));
 
     this.server = new Server(url);
 
@@ -60,6 +62,37 @@ export default class WS {
     });
   }
 
+  get connected() {
+    let resolve: (socket: WebSocket) => void;
+    const connectedPromise = new Promise<WebSocket>((done) => (resolve = done));
+    const waitForConnected = async () => {
+      await act(async () => {
+        await this._isConnected;
+      });
+      resolve(await this._isConnected); // make sure `await act` is really done
+    };
+    waitForConnected();
+    return connectedPromise;
+  }
+
+  get closed() {
+    let resolve: () => void;
+    const closedPromise = new Promise<void>((done) => (resolve = done));
+    const waitForclosed = async () => {
+      await act(async () => {
+        await this._isClosed;
+      });
+      await this._isClosed; // make sure `await act` is really done
+      resolve();
+    };
+    waitForclosed();
+    return closedPromise;
+  }
+
+  get nextMessage() {
+    return this.messagesToConsume.get();
+  }
+
   on(
     eventName: "connection" | "message" | "close",
     callback: (socket: MockWebSocket) => void
@@ -68,20 +101,22 @@ export default class WS {
     this.server.on(eventName, callback);
   }
 
-  get nextMessage() {
-    return this.messagesToConsume.get();
-  }
-
   send(message: DeserializedMessage) {
-    this.server.emit("message", this.serializer(message));
+    act(() => {
+      this.server.emit("message", this.serializer(message));
+    });
   }
 
   close(options?: CloseOptions) {
-    this.server.close(options);
+    act(() => {
+      this.server.close(options);
+    });
   }
 
   error() {
-    this.server.emit("error", null);
+    act(() => {
+      this.server.emit("error", null);
+    });
     this.server.close();
   }
 }