Integration tests for Similarity Search API.

tools/similarity_search/package.json

@@ -7,7 +7,7 @@
   "dependencies": {
     "@cloudflare/vitest-pool-workers": "^0.2.6",
     "hono": "^4.3.2",
-    "vitest": "1.3.0"
+    "vitest": "1.5.0"
   },
   "devDependencies": {
     "@cloudflare/workers-types": "^4.20240403.0",

tools/similarity_search/src/index.ts

@@ -1,6 +1,6 @@
 import { Hono } from "hono"
 
-type Env = {
+export type Env = {
   API_KEY_TOKEN_CHECK: string
   AI: Ai
   VECTORIZE_INDEX: VectorizeIndex

tools/similarity_search/test/env.d.ts

@@ -0,0 +1,5 @@
+import { Env } from "../src/index"
+
+declare module "cloudflare:test" {
+  interface ProvidedEnv extends Env {}
+}

tools/similarity_search/test/index.spec.ts

@@ -1,14 +1,19 @@
-import { SELF } from "cloudflare:test"
-import { describe, it, expect } from "vitest"
+import { SELF, env } from "cloudflare:test"
+import { describe, it, expect, vi } from "vitest"
 
 import "../src/index"
 
+interface SimilarityCheckResponse {
+  similarity_score: number
+}
+
 describe("Authentication", () => {
-  it("returns 401 Unauthorized when API key is missing or invalid", async () => {
+  it("returns 401 when API key is invalid", async () => {
     const response = await SELF.fetch("https://example.com/", {
       method: "POST",
       headers: {
-        "Content-Type": "application/json"
+        "Content-Type": "application/json",
+        "X-API-Key": "invalid-api-key"
       },
       body: JSON.stringify({
         text: "Sample text",
@@ -17,6 +22,91 @@ describe("Authentication", () => {
     })
 
     expect(response.status).toBe(401)
-    expect(await response.text()).toBe("Unauthorized")
+  })
+})
+
+describe("Validation", () => {
+  it.each([
+    ["text is missing", { namespace: "test-namespace" }],
+    ["namespace is missing", { text: "Sample text" }]
+  ])("returns 400 when %s", async (description, body) => {
+    const response = await SELF.fetch("https://example.com/", {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        "X-API-Key": "test-api-key"
+      },
+      body: JSON.stringify(body)
+    })
+
+    expect(response.status).toBe(400)
+  })
+})
+
+describe("Functionality", () => {
+  it("queries AI model during main route call and gets proper response", async () => {
+    const AICall = vi.spyOn(env.AI, "run")
+
+    const response = await SELF.fetch("https://example.com/", {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        "X-API-Key": "test-api-key"
+      },
+      body: JSON.stringify({
+        text: "Sample text",
+        namespace: "test-namespace"
+      })
+    })
+
+    expect(AICall).toHaveBeenCalledWith("@cf/baai/bge-base-en-v1.5", {
+      text: ["Sample text"]
+    })
+    expect(AICall).toHaveReturnedWith({ data: {} })
+
+    expect(response.status).toBe(200)
+  })
+
+  it("queries vector database during main route call and gets proper response", async () => {
+    const queryCall = vi.spyOn(env.VECTORIZE_INDEX, "query")
+
+    const response = await SELF.fetch("https://example.com/", {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        "X-API-Key": "test-api-key"
+      },
+      body: JSON.stringify({
+        text: "Sample text",
+        namespace: "test-namespace"
+      })
+    })
+
+    expect(queryCall).toHaveBeenCalledWith(undefined, {
+      namespace: "test-namespace",
+      topK: 1
+    })
+    expect(queryCall).toHaveReturnedWith({ matches: [{ score: 0.5678 }] })
+
+    expect(response.status).toBe(200)
+  })
+
+  it("returns similarity score for valid requests", async () => {
+    const response = await SELF.fetch("https://example.com/", {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        "X-API-Key": "test-api-key",
+      },
+      body: JSON.stringify({
+        text: "Sample text",
+        namespace: "test-namespace",
+      }),
+    })
+
+    expect(response.status).toBe(200)
+    const jsonResponse: SimilarityCheckResponse = await response.json()
+    expect(jsonResponse).toHaveProperty("similarity_score")
+    expect(typeof jsonResponse.similarity_score).toBe("number")
   })
 })