add vitest fixture that documents AI/Vectorize bindings with vitest (#8624)

* add ai fixture

* Update vitest.config.ts

* Update fixtures/vitest-pool-workers-examples/ai-vectorize/global-setup.ts

Co-authored-by: Carmen Popoviciu <[email protected]>

* add config warning

* rename worker

* lint

* changeset

* Update config.ts

Co-authored-by: Carmen Popoviciu <[email protected]>

* Update packages/vitest-pool-workers/src/pool/config.ts

* fix

---------

Co-authored-by: Carmen Popoviciu <[email protected]>

Author: emily-shen

.changeset/warm-poets-feel.md

@@ -0,0 +1,7 @@
+---
+"@cloudflare/vitest-pool-workers": patch
+---
+
+fix: Add usage charge warning when Vectorize and AI bindings are used in Vitest
+
+Vectorize and AI bindings can now be used with Vitest. However, because they have no local simulators, they will access your account and incur usage charges, even in testing. Therefore we recommend mocking any usage of these bindings when testing.

fixtures/vitest-pool-workers-examples/ai-vectorize/README.md

@@ -0,0 +1,7 @@
+# 🤖 AI and Vectorize
+
+This Worker uses the AI and Vectorize bindings. @cloudflare/vitest-pool-workers@^0.8.1 is required to use AI and Vectorize bindings in the Vitest integration.
+
+[!WARNING]
+
+Because Workers AI and Vectorize bindings do not have a local simulator, usage of these bindings will always access your Cloudflare account, and so will incur usage charges even in local development and testing. We recommend mocking any usage of these bindings in your tests as demonstrated in this fixture.

fixtures/vitest-pool-workers-examples/ai-vectorize/global-setup.ts

@@ -0,0 +1,9 @@
+import childProcess from "node:child_process";
+
+// Global setup runs inside Node.js, not `workerd`
+export default function () {
+	const label = "Built ai-vectorize Worker";
+	console.time(label);
+	childProcess.execSync("wrangler build", { cwd: __dirname });
+	console.timeEnd(label);
+}

fixtures/vitest-pool-workers-examples/ai-vectorize/src/env.d.ts

@@ -0,0 +1,8 @@
+// Generated by Wrangler by running `wrangler types` (hash: 266f7400cc06616d9ccb541889857ea7)
+declare namespace Cloudflare {
+	interface Env {
+		VECTORIZE: VectorizeIndex;
+		AI: Ai;
+	}
+}
+interface Env extends Cloudflare.Env {}

fixtures/vitest-pool-workers-examples/ai-vectorize/src/index.ts

@@ -0,0 +1,36 @@
+interface EmbeddingResponse {
+	shape: number[];
+	data: number[][];
+}
+
+export default {
+	async fetch(request, env, ctx): Promise<Response> {
+		const path = new URL(request.url).pathname;
+
+		if (path === "/ai") {
+			const stories = [
+				"This is a story about an orange cloud",
+				"This is a story about a llama",
+			];
+			const modelResp: EmbeddingResponse = await env.AI.run(
+				"@cf/baai/bge-base-en-v1.5",
+				{
+					text: stories,
+				}
+			);
+			return Response.json(modelResp);
+		}
+
+		if (path === "/vectorize") {
+			const vectors: VectorizeVector[] = [
+				{ id: "123", values: [...Array(768).keys()] },
+				{ id: "456", values: [...Array(768).keys()] },
+			];
+
+			let inserted = await env.VECTORIZE.upsert(vectors);
+			return Response.json(inserted);
+		}
+
+		return new Response("Hello World!");
+	},
+} satisfies ExportedHandler<Env>;

fixtures/vitest-pool-workers-examples/ai-vectorize/src/tsconfig.json

@@ -0,0 +1,4 @@
+{
+	"extends": "../../tsconfig.workerd.json",
+	"include": ["./**/*.ts"]
+}

fixtures/vitest-pool-workers-examples/ai-vectorize/test/env.d.ts

@@ -0,0 +1,3 @@
+declare module "cloudflare:test" {
+	interface ProvidedEnv extends Env {}
+}

fixtures/vitest-pool-workers-examples/ai-vectorize/test/index.spec.ts

@@ -0,0 +1,76 @@
+import {
+	createExecutionContext,
+	env,
+	SELF,
+	waitOnExecutionContext,
+} from "cloudflare:test";
+import { describe, expect, it, vi } from "vitest";
+import worker from "../src/index";
+
+// For now, you'll need to do something like this to get a correctly-typed
+// `Request` to pass to `worker.fetch()`.
+const IncomingRequest = Request<unknown, IncomingRequestCfProperties>;
+
+describe("Tests that do hit the AI binding", () => {
+	describe("unit style", () => {
+		it("lets you mock the ai binding", async () => {
+			const request = new IncomingRequest("http://example.com/ai");
+
+			const ctx = createExecutionContext();
+
+			// mock the AI run function by directly modifying `env`
+			env.AI.run = vi.fn().mockResolvedValue({
+				shape: [1, 2],
+				data: [[0, 0]],
+			});
+			const response = await worker.fetch(request, env, ctx);
+
+			await waitOnExecutionContext(ctx);
+			expect(await response.text()).toMatchInlineSnapshot(
+				`"{"shape":[1,2],"data":[[0,0]]}"`
+			);
+		});
+
+		it("lets you mock the vectorize binding", async () => {
+			const request = new IncomingRequest("http://example.com/vectorize");
+			const ctx = createExecutionContext();
+
+			// mock the vectorize upsert function by directly modifying `env`
+			const mockVectorizeStore: VectorizeVector[] = [];
+			env.VECTORIZE.upsert = vi.fn().mockImplementation(async (vectors) => {
+				mockVectorizeStore.push(...vectors);
+				return { mutationId: "123" };
+			});
+
+			const response = await worker.fetch(request, env, ctx);
+
+			await waitOnExecutionContext(ctx);
+			expect(await response.text()).toMatchInlineSnapshot(
+				`"{"mutationId":"123"}"`
+			);
+			expect(mockVectorizeStore.map((v) => v.id)).toMatchInlineSnapshot(`
+				[
+				  "123",
+				  "456",
+				]
+			`);
+		});
+	});
+});
+
+describe("Tests that do not hit the AI binding", () => {
+	it("responds with Hello World! (unit style)", async () => {
+		const request = new IncomingRequest("http://example.com");
+		// Create an empty context to pass to `worker.fetch()`.
+		const ctx = createExecutionContext();
+		const response = await worker.fetch(request, env, ctx);
+		// Wait for all `Promise`s passed to `ctx.waitUntil()` to settle before running test assertions
+		await waitOnExecutionContext(ctx);
+		expect(await response.text()).toMatchInlineSnapshot(`"Hello World!"`);
+	});
+
+	it("responds with Hello World! (integration style)", async () => {
+		const response = await SELF.fetch("https://example.com");
+		expect(await response.text()).toMatchInlineSnapshot(`"Hello World!"`);
+	});
+});

fixtures/vitest-pool-workers-examples/ai-vectorize/test/tsconfig.json

@@ -0,0 +1,4 @@
+{
+	"extends": "../../tsconfig.workerd-test.json",
+	"include": ["./**/*.ts", "../src/env.d.ts"]
+}

fixtures/vitest-pool-workers-examples/ai-vectorize/tsconfig.json

@@ -0,0 +1,4 @@
+{
+	"extends": "../tsconfig.node.json",
+	"include": ["./*.ts"]
+}