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,22 +1,98 @@
-import { SELF } from "cloudflare:test"
+import { SELF, env } from "cloudflare:test"
 import { describe, it, expect } 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 Unauthorized 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",
-        namespace: "test-namespace"
-      })
+        namespace: "test-namespace",
+      }),
     })
 
     expect(response.status).toBe(401)
     expect(await response.text()).toBe("Unauthorized")
   })
 })
+
+describe("Validation", () => {
+  it("returns 400 Invalid JSON format when text is missing", 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({
+        namespace: "test-namespace",
+      }),
+    })
+
+    expect(response.status).toBe(400)
+    expect(await response.text()).toBe("Invalid JSON format")
+  })
+
+  it("returns 400 Invalid JSON format when namespace is missing", 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",
+      }),
+    })
+
+    expect(response.status).toBe(400)
+    expect(await response.text()).toBe("Invalid JSON format")
+  })
+})
+
+describe("Functionality", () => {
+  it("runs AI model and gets vectorized string back", async () => {
+    const response = await env.AI.run("@cf/baai/bge-base-en-v1.5", {
+      text: ["Sample text"],
+    })
+    expect(response).toHaveProperty("data")
+  })
+
+  it("queries vector database and gets proper response", async () => {
+    const response = await env.VECTORIZE_INDEX.query([1, 2, 3], {
+      namespace: "test-namespace",
+      topK: 1,
+    })
+    expect(response).toHaveProperty("matches")
+    expect(response.matches.length).toBeGreaterThan(0)
+  })
+
+  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")
+  })
+})