Feature: getMiniflareWaitUntil (#345)

* getMiniflareWaitUntil addded

* tests pass

* added ExecutionContext; WaitUntil Promises; resolve immediately

* getMiniflareWaitUntil addded

* set ExecutionContext to global; adjust tests accordingly

* getMiniflareWaitUntil addded

* cleanup; fix ava

* remove extra waitUntil file

* Add `getMiniflareWaitUntil()` support to Vitest environment

Co-authored-by: CraigglesO <[email protected]>

Co-authored-by: bcoll <[email protected]>

Author: CraigglesO

packages/core/src/standards/event.ts

@@ -34,7 +34,7 @@ const SUGGEST_GLOBAL_BINDING_MODULES =
 
 const kResponse = Symbol("kResponse");
 const kPassThrough = Symbol("kPassThrough");
-const kWaitUntil = Symbol("kWaitUntil");
+export const kWaitUntil = Symbol("kWaitUntil");
 const kSent = Symbol("kSent");
 
 export class FetchEvent extends Event {

packages/jest-environment-miniflare/src/index.ts

@@ -10,6 +10,7 @@ import { MiniflareCore } from "@miniflare/core";
 import { QueueBroker } from "@miniflare/queues";
 import { VMScriptRunner, defineHasInstances } from "@miniflare/runner-vm";
 import {
+  ExecutionContext,
   PLUGINS,
   StackedMemoryStorageFactory,
   createMiniflareEnvironment,
@@ -127,6 +128,7 @@ export default class MiniflareEnvironment implements JestEnvironment<Timer> {
       },
       this.config.testEnvironmentOptions,
       {
+        ExecutionContext,
         // Make sure fancy jest console and faked timers are included
         console: global.console,
         setTimeout: global.setTimeout,

packages/jest-environment-miniflare/test/fixtures/waitUntil.worker.spec.js

@@ -0,0 +1,73 @@
+function testResponse(body) {
+  return new Response(body, { headers: { "Cache-Control": "max-age=3600" } });
+}
+
+test("FetchEvent: waitUntil: await resolving promise array using getMiniflareWaitUntil", async () => {
+  // build a FetchEvent:
+  const url = "http://localhost/fetchEvent/waitUntil";
+  const request = new Request(url);
+  const fetchEvent = new FetchEvent("fetch", { request });
+
+  // run a waitUntil
+  fetchEvent.waitUntil(caches.default.put(url, testResponse("example cache")));
+
+  // ensure that waitUntil has yet to be run
+  let cachedResponse = await caches.default.match(url);
+  expect(cachedResponse).toBeUndefined();
+
+  // pull in the waitUntil stack
+  const waitUntilList = await getMiniflareWaitUntil(fetchEvent);
+  expect(waitUntilList).not.toBeNull();
+
+  cachedResponse = await caches.default.match(url);
+  expect(cachedResponse).toBeTruthy();
+  expect(await cachedResponse.text()).toBe("example cache");
+});
+
+test("ScheduledEvent: waitUntil: await resolving promise array using getMiniflareWaitUntil", async () => {
+  const url = "http://localhost/scheduledEvent/waitUntil";
+
+  const scheduledEvent = new ScheduledEvent("scheduled", {
+    scheduledTime: 1000,
+    cron: "30 * * * *",
+  });
+
+  // run a waitUntil
+  scheduledEvent.waitUntil(
+    caches.default.put(url, testResponse("example cache"))
+  );
+
+  // ensure that waitUntil has yet to be run
+  let cachedResponse = await caches.default.match(url);
+  expect(cachedResponse).toBeUndefined();
+
+  // pull in the waitUntil stack
+  const waitUntilList = await getMiniflareWaitUntil(scheduledEvent);
+  expect(waitUntilList).not.toBeNull();
+
+  cachedResponse = await caches.default.match(url);
+  expect(cachedResponse).toBeTruthy();
+  expect(await cachedResponse.text()).toBe("example cache");
+});
+
+test("ExecutionContext: waitUntil: await resolving promise array using getMiniflareWaitUntil", async () => {
+  const url = "http://localhost/ctx/waitUntil";
+
+  const ctx = new ExecutionContext();
+  expect(Object.getPrototypeOf(ctx).constructor.name).toBe("ExecutionContext");
+
+  // run a waitUntil
+  ctx.waitUntil(caches.default.put(url, testResponse("example cache")));
+
+  // ensure that waitUntil has yet to be run
+  let cachedResponse = await caches.default.match(url);
+  expect(cachedResponse).toBeUndefined();
+
+  // pull in the waitUntil stack
+  const waitUntilList = await getMiniflareWaitUntil(ctx);
+  expect(waitUntilList).not.toBeNull();
+
+  cachedResponse = await caches.default.match(url);
+  expect(cachedResponse).toBeTruthy();
+  expect(await cachedResponse.text()).toBe("example cache");
+});

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

@@ -1,4 +1,9 @@
-import { MiniflareCore } from "@miniflare/core";
+import {
+  FetchEvent,
+  MiniflareCore,
+  ScheduledEvent,
+  kWaitUntil,
+} from "@miniflare/core";
 import {
   DurableObjectId,
   DurableObjectStorage,
@@ -7,12 +12,25 @@ import { Context } from "@miniflare/shared";
 import { MockAgent } from "undici";
 import { PLUGINS } from "./plugins";
 
+export class ExecutionContext {
+  [kWaitUntil]: Promise<unknown>[] = [];
+
+  passThroughOnException(): void {}
+
+  waitUntil(promise: Promise<any>): void {
+    this[kWaitUntil].push(promise);
+  }
+}
+
 declare global {
   function getMiniflareBindings<Bindings = Context>(): Bindings;
   function getMiniflareDurableObjectStorage(
     id: DurableObjectId
   ): Promise<DurableObjectStorage>;
   function getMiniflareFetchMock(): MockAgent;
+  function getMiniflareWaitUntil<WaitUntil extends any[] = unknown[]>(
+    event: FetchEvent | ScheduledEvent | ExecutionContext
+  ): Promise<WaitUntil>;
   function flushMiniflareDurableObjectAlarms(
     ids: DurableObjectId[]
   ): Promise<void>;
@@ -24,6 +42,9 @@ export interface MiniflareEnvironmentUtilities {
     id: DurableObjectId
   ): Promise<DurableObjectStorage>;
   getMiniflareFetchMock(): MockAgent;
+  getMiniflareWaitUntil<WaitUntil extends any[] = unknown[]>(
+    event: FetchEvent | ScheduledEvent | ExecutionContext
+  ): Promise<WaitUntil>;
   flushMiniflareDurableObjectAlarms(ids: DurableObjectId[]): Promise<void>;
 }
 
@@ -48,6 +69,11 @@ export async function createMiniflareEnvironmentUtilities(
     getMiniflareFetchMock() {
       return fetchMock;
     },
+    getMiniflareWaitUntil<WaitUntil extends any[] = unknown[]>(
+      event: FetchEvent | ScheduledEvent | ExecutionContext
+    ): Promise<WaitUntil> {
+      return Promise.all(event[kWaitUntil]) as Promise<WaitUntil>;
+    },
     async flushMiniflareDurableObjectAlarms(ids?: DurableObjectId[]) {
       const plugin = (await mf.getPlugins()).DurableObjectsPlugin;
       const storage = mf.getPluginStorage("DurableObjectsPlugin");

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

@@ -10,6 +10,7 @@ import { PLUGINS } from "./plugins";
 
 export * from "./plugins";
 export * from "./storage";
+export { ExecutionContext } from "./globals";
 export type { MiniflareEnvironmentUtilities } from "./globals";
 
 const log = new NoOpLog();

packages/vitest-environment-miniflare/src/index.ts

@@ -1,6 +1,7 @@
 import { QueueBroker } from "@miniflare/queues";
 import { VMScriptRunner } from "@miniflare/runner-vm";
 import {
+  ExecutionContext,
   StackedMemoryStorageFactory,
   createMiniflareEnvironment,
 } from "@miniflare/shared-test-environment";
@@ -66,7 +67,8 @@ export default <Environment>{
     const storageFactory = new StackedMemoryStorageFactory();
     const [mf, mfGlobalScope] = await createMiniflareEnvironment(
       { storageFactory, scriptRunner, queueBroker },
-      options
+      options,
+      { ExecutionContext }
     );
 
     // Attach isolated storage setup function

packages/vitest-environment-miniflare/test/fixtures/service-worker/waitUntil.worker.spec.js

@@ -0,0 +1,76 @@
+import { expect, test } from "vitest";
+setupMiniflareIsolatedStorage();
+
+function testResponse(body) {
+  return new Response(body, { headers: { "Cache-Control": "max-age=3600" } });
+}
+
+test("FetchEvent: waitUntil: await resolving promise array using getMiniflareWaitUntil", async () => {
+  // build a FetchEvent:
+  const url = "http://localhost/fetchEvent/waitUntil";
+  const request = new Request(url);
+  const fetchEvent = new FetchEvent("fetch", { request });
+
+  // run a waitUntil
+  fetchEvent.waitUntil(caches.default.put(url, testResponse("example cache")));
+
+  // ensure that waitUntil has yet to be run
+  let cachedResponse = await caches.default.match(url);
+  expect(cachedResponse).toBeUndefined();
+
+  // pull in the waitUntil stack
+  const waitUntilList = await getMiniflareWaitUntil(fetchEvent);
+  expect(waitUntilList).not.toBeNull();
+
+  cachedResponse = await caches.default.match(url);
+  expect(cachedResponse).toBeTruthy();
+  expect(await cachedResponse.text()).toBe("example cache");
+});
+
+test("ScheduledEvent: waitUntil: await resolving promise array using getMiniflareWaitUntil", async () => {
+  const url = "http://localhost/scheduledEvent/waitUntil";
+
+  const scheduledEvent = new ScheduledEvent("scheduled", {
+    scheduledTime: 1000,
+    cron: "30 * * * *",
+  });
+
+  // run a waitUntil
+  scheduledEvent.waitUntil(
+    caches.default.put(url, testResponse("example cache"))
+  );
+
+  // ensure that waitUntil has yet to be run
+  let cachedResponse = await caches.default.match(url);
+  expect(cachedResponse).toBeUndefined();
+
+  // pull in the waitUntil stack
+  const waitUntilList = await getMiniflareWaitUntil(scheduledEvent);
+  expect(waitUntilList).not.toBeNull();
+
+  cachedResponse = await caches.default.match(url);
+  expect(cachedResponse).toBeTruthy();
+  expect(await cachedResponse.text()).toBe("example cache");
+});
+
+test("ExecutionContext: waitUntil: await resolving promise array using getMiniflareWaitUntil", async () => {
+  const url = "http://localhost/ctx/waitUntil";
+
+  const ctx = new ExecutionContext();
+  expect(Object.getPrototypeOf(ctx).constructor.name).toBe("ExecutionContext");
+
+  // run a waitUntil
+  ctx.waitUntil(caches.default.put(url, testResponse("example cache")));
+
+  // ensure that waitUntil has yet to be run
+  let cachedResponse = await caches.default.match(url);
+  expect(cachedResponse).toBeUndefined();
+
+  // pull in the waitUntil stack
+  const waitUntilList = await getMiniflareWaitUntil(ctx);
+  expect(waitUntilList).not.toBeNull();
+
+  cachedResponse = await caches.default.match(url);
+  expect(cachedResponse).toBeTruthy();
+  expect(await cachedResponse.text()).toBe("example cache");
+});