Add runWithMiniflareDurableObjectGates for test environments

...allowing you to use I/O gates in the Jest/Vitest environments

Author: mrbbot

packages/durable-objects/src/index.ts

@@ -5,6 +5,7 @@ export {
   DurableObjectState,
   DurableObjectStub,
   DurableObjectNamespace,
+  _kRunWithGates,
 } from "./namespace";
 export type {
   DurableObject,

packages/durable-objects/src/namespace.ts

@@ -62,6 +62,8 @@ export interface DurableObject {
 }
 
 export const kInstance = Symbol("kInstance");
+/** @internal */
+export const _kRunWithGates = Symbol("kRunWithGates");
 export const kAlarm = Symbol("kAlarm");
 const kFetch = Symbol("kFetch");
 
@@ -81,41 +83,39 @@ export class DurableObjectState {
     return this.#inputGate.runWithClosed(closure);
   }
 
-  [kFetch](request: Request): Promise<Response> {
+  [_kRunWithGates]<T>(closure: () => Awaitable<T>): Promise<T> {
     // TODO: catch, reset object on error
     const outputGate = new OutputGate();
-    return outputGate.runWith(() =>
-      this.#inputGate.runWith(() => {
-        const instance = this[kInstance];
-        if (!instance?.fetch) {
-          throw new DurableObjectError(
-            "ERR_NO_HANDLER",
-            "No fetch handler defined in Durable Object"
-          );
-        }
-        return instance.fetch(request);
-      })
-    );
+    return outputGate.runWith(() => this.#inputGate.runWith(closure));
+  }
+
+  [kFetch](request: Request): Promise<Response> {
+    return this[_kRunWithGates](() => {
+      const instance = this[kInstance];
+      if (!instance?.fetch) {
+        throw new DurableObjectError(
+          "ERR_NO_HANDLER",
+          "No fetch handler defined in Durable Object"
+        );
+      }
+      return instance.fetch(request);
+    });
   }
 
   [kAlarm](): Promise<void> {
-    // TODO: catch, reset object on error
-    const outputGate = new OutputGate();
-    return outputGate.runWith(() =>
-      this.#inputGate.runWith(async () => {
-        // Delete the local alarm
-        await this.storage.deleteAlarm();
-        // Grab the instance and call the alarm handler
-        const instance = this[kInstance];
-        if (!instance?.alarm) {
-          throw new DurableObjectError(
-            "ERR_NO_HANDLER",
-            "No alarm handler defined in Durable Object"
-          );
-        }
-        return instance.alarm();
-      })
-    );
+    return this[_kRunWithGates](async () => {
+      // Delete the local alarm
+      await this.storage.deleteAlarm();
+      // Grab the instance and call the alarm handler
+      const instance = this[kInstance];
+      if (!instance?.alarm) {
+        throw new DurableObjectError(
+          "ERR_NO_HANDLER",
+          "No alarm handler defined in Durable Object"
+        );
+      }
+      return instance.alarm();
+    });
   }
 }
 

packages/jest-environment-miniflare/test/fixtures/durableobjects.module.spec.js

@@ -1,5 +1,3 @@
-import { TestObject } from "./module-worker.js";
-
 beforeAll(async () => {
   const { TEST_OBJECT } = getMiniflareBindings();
   const id = TEST_OBJECT.idFromName("test");
@@ -16,11 +14,32 @@ test("Durable Objects", async () => {
 });
 
 test("Durable Objects direct", async () => {
+  class Counter {
+    constructor(state) {
+      this.storage = state.storage;
+    }
+    async fetch() {
+      const count = (await this.storage.get("count")) ?? 0;
+      void this.storage.put("count", count + 1);
+      return new Response(String(count));
+    }
+  }
+
   // https://github.com/cloudflare/miniflare/issues/157
   const env = getMiniflareBindings();
+  // Doesn't matter too much that we're using a different object binding here
   const id = env.TEST_OBJECT.idFromName("test");
   const state = await getMiniflareDurableObjectState(id);
-  const object = new TestObject(state, env);
-  const res = await object.fetch(new Request("https://object/"));
-  expect(await res.text()).toBe("durable:https://object/:value");
+  const object = new Counter(state, env);
+  const [res1, res2] = await Promise.all([
+    runWithMiniflareDurableObjectGates(state, () =>
+      object.fetch(new Request("https://object/"))
+    ),
+    runWithMiniflareDurableObjectGates(state, () =>
+      object.fetch(new Request("https://object/"))
+    ),
+  ]);
+  expect(await state.storage.get("count")).toBe(2);
+  expect(await res1.text()).toBe("0");
+  expect(await res2.text()).toBe("1");
 });

packages/shared-test-environment/src/globals.ts

@@ -8,8 +8,9 @@ import {
   DurableObjectId,
   DurableObjectState,
   DurableObjectStorage,
+  _kRunWithGates,
 } from "@miniflare/durable-objects";
-import { Context } from "@miniflare/shared";
+import { Awaitable, Context } from "@miniflare/shared";
 import { MockAgent } from "undici";
 import { PLUGINS } from "./plugins";
 
@@ -31,6 +32,10 @@ declare global {
   function getMiniflareDurableObjectState(
     id: DurableObjectId
   ): Promise<DurableObjectState>;
+  function runWithMiniflareDurableObjectGates<T>(
+    state: DurableObjectState,
+    closure: () => Awaitable<T>
+  ): Promise<T>;
   function getMiniflareFetchMock(): MockAgent;
   function getMiniflareWaitUntil<WaitUntil extends any[] = unknown[]>(
     event: FetchEvent | ScheduledEvent | ExecutionContext
@@ -48,6 +53,10 @@ export interface MiniflareEnvironmentUtilities {
   getMiniflareDurableObjectState(
     id: DurableObjectId
   ): Promise<DurableObjectState>;
+  runWithMiniflareDurableObjectGates<T>(
+    state: DurableObjectState,
+    closure: () => Awaitable<T>
+  ): Promise<T>;
   getMiniflareFetchMock(): MockAgent;
   getMiniflareWaitUntil<WaitUntil extends any[] = unknown[]>(
     event: FetchEvent | ScheduledEvent | ExecutionContext
@@ -78,6 +87,12 @@ export async function createMiniflareEnvironmentUtilities(
       const storage = plugin.getStorage(factory, id);
       return new DurableObjectState(id, storage);
     },
+    runWithMiniflareDurableObjectGates<T>(
+      state: DurableObjectState,
+      closure: () => Awaitable<T>
+    ) {
+      return state[_kRunWithGates](closure);
+    },
     getMiniflareFetchMock() {
       return fetchMock;
     },

packages/vitest-environment-miniflare/test/fixtures/modules/durableobjects.module.spec.js

@@ -1,5 +1,4 @@
 import { beforeAll, expect, test } from "vitest";
-import { TestObject } from "./module-worker.js";
 setupMiniflareIsolatedStorage();
 
 beforeAll(async () => {
@@ -18,11 +17,32 @@ test("Durable Objects", async () => {
 });
 
 test("Durable Objects direct", async () => {
+  class Counter {
+    constructor(state) {
+      this.storage = state.storage;
+    }
+    async fetch() {
+      const count = (await this.storage.get("count")) ?? 0;
+      void this.storage.put("count", count + 1);
+      return new Response(String(count));
+    }
+  }
+
   // https://github.com/cloudflare/miniflare/issues/157
   const env = getMiniflareBindings();
+  // Doesn't matter too much that we're using a different object binding here
   const id = env.TEST_OBJECT.idFromName("test");
   const state = await getMiniflareDurableObjectState(id);
-  const object = new TestObject(state, env);
-  const res = await object.fetch(new Request("https://object/"));
-  expect(await res.text()).toBe("durable:https://object/:value");
+  const object = new Counter(state, env);
+  const [res1, res2] = await Promise.all([
+    runWithMiniflareDurableObjectGates(state, () =>
+      object.fetch(new Request("https://object/"))
+    ),
+    runWithMiniflareDurableObjectGates(state, () =>
+      object.fetch(new Request("https://object/"))
+    ),
+  ]);
+  expect(await state.storage.get("count")).toBe(2);
+  expect(await res1.text()).toBe("0");
+  expect(await res2.text()).toBe("1");
 });