From 8d93ebcf037726fcc4cabf06caf89a252ae397d8 Mon Sep 17 00:00:00 2001
From: Gabriel Massadas <gmassadas@cloudflare.com>
Date: Thu, 21 Nov 2024 19:27:58 +0000
Subject: [PATCH] Lint files

---
 .../wrangler/src/__tests__/ai.local.test.ts   | 68 +++++++++++++------
 1 file changed, 48 insertions(+), 20 deletions(-)

diff --git a/packages/wrangler/src/__tests__/ai.local.test.ts b/packages/wrangler/src/__tests__/ai.local.test.ts
index 7b11bf8eb8bb..1949e558391f 100644
--- a/packages/wrangler/src/__tests__/ai.local.test.ts
+++ b/packages/wrangler/src/__tests__/ai.local.test.ts
@@ -1,41 +1,69 @@
 import { Request } from "miniflare";
-import { http, HttpResponse } from "msw";
+import { HttpResponse } from "msw";
 import { AIFetcher } from "../ai/fetcher";
-import { msw } from "./helpers/msw";
+import * as internal from "../cfetch/internal";
+import * as user from "../user";
+import type { RequestInit } from "undici";
 
 describe("ai", () => {
 	describe("fetcher", () => {
+		afterEach(() => {
+			vi.restoreAllMocks();
+		});
+
 		describe("local", () => {
 			it("should send x-forwarded-host header", async () => {
-				mockAIProxyRequest();
+				vi.spyOn(user, "getAccountId").mockImplementation(async () => "123");
+				vi.spyOn(internal, "performApiFetch").mockImplementation(
+					async (resource: string, init: RequestInit = {}) => {
+						const headers = new Headers(init.headers);
+						return HttpResponse.json({
+							xForwardedFor: headers.get("X-Forwarded-Host"),
+							method: init.method,
+						});
+					}
+				);
+
+				const url = "http://internal.ai/ai/test/path?version=123";
+				const resp = await AIFetcher(
+					new Request(url, {
+						method: "PATCH",
+						headers: {
+							"x-example": "test",
+						},
+					})
+				);
+
+				expect(await resp.json()).toEqual({
+					xForwardedFor: url,
+					method: "PATCH",
+				});
+			});
+
+			it("account id should be set", async () => {
+				vi.spyOn(user, "getAccountId").mockImplementation(async () => "123");
+				vi.spyOn(internal, "performApiFetch").mockImplementation(
+					async (resource: string) => {
+						return HttpResponse.json({
+							resource: resource,
+						});
+					}
+				);
 
 				const url = "http://internal.ai/ai/test/path?version=123";
 				const resp = await AIFetcher(
 					new Request(url, {
-						method: 'PATCH',
+						method: "PATCH",
 						headers: {
 							"x-example": "test",
 						},
 					})
 				);
 
-				expect(await resp.json()).toEqual({ xForwardedFor: url, method: 'PATCH' });
+				expect(await resp.json()).toEqual({
+					resource: "/accounts/123/ai/run/proxy",
+				});
 			});
 		});
 	});
 });
-
-function mockAIProxyRequest() {
-	msw.use(
-		http.get(
-			"*/accounts/:accountId/ai/run/proxy",
-			(c) => {
-				return HttpResponse.json({
-					xForwardedFor: c.request.headers.get("X-Forwarded-Host"),
-					method: c.request.method
-				});
-			},
-			{ once: true }
-		)
-	);
-}