add test for stream cache

Signed-off-by: Mengxin Liu <[email protected]>
oilbeater · Sep 22, 2024 · b266f53 · b266f53
1 parent 5661a1b
commit b266f53
Showing 1 changed file with 53 additions and 30 deletions.
diff --git a/test/index.spec.ts b/test/index.spec.ts
@@ -20,37 +20,38 @@ describe('Welcome to Malacca worker', () => {
 });
 
 describe('Test Cache', () => {
-	console.log(import.meta.env)
 	const url = `https://example.com/azure-openai/${import.meta.env.VITE_AZURE_RESOURCE_NAME}/${import.meta.env.VITE_AZURE_DEPLOYMENT_NAME}/chat/completions?api-version=2024-07-01-preview`;
-	const body = `
-	{
-  "messages": [
-    {
-      "role": "system",
-      "content": [
-        {
-          "type": "text",
-          "text": "You are an AI assistant that helps people find information."
-        }
-      ]
-    },
-    {
-      "role": "user",
-      "content": [
-        {
-          "type": "text",
-          "text": "Tell me a very short story about sunwukong"
-        }
-      ]
-    }
-  ],
-  "temperature": 0.7,
-  "top_p": 0.95,
-  "max_tokens": 800
-}`
-	console.log(url)
+  const createRequestBody = (stream: boolean) => `
+  {
+    "messages": [
+      {
+        "role": "system",
+        "content": [
+          {
+            "type": "text",
+            "text": "You are an AI assistant that helps people find information."
+          }
+        ]
+      },
+      {
+        "role": "user",
+        "content": [
+          {
+            "type": "text",
+            "text": "Tell me a very short story about sunwukong"
+          }
+        ]
+      }
+    ],
+    "temperature": 0.7,
+    "top_p": 0.95,
+    "max_tokens": 800,
+    "stream": ${stream}
+  }`;
+
 	it('with cache first response should with header malacca-cache-status: miss and following response with hit', async () => {
-		let start = Date.now();
+		const body = createRequestBody(false);
+    let start = Date.now();
     let response = await SELF.fetch(url, { method: 'POST', body: body, headers: { 'Content-Type': 'application/json', 'api-key': 'oilbeater' } });
     const value = await response.json()
     const duration = Date.now() - start
@@ -67,5 +68,27 @@ describe('Test Cache', () => {
     expect(response.headers.get('malacca-cache-status')).toBe('hit');
     expect(value).toEqual(cacheValue)
     expect(duration/2).toBeGreaterThan(cacheDuration)
-	});
+  });
+
+  it('Test stream with cache', async () => { 
+    const body = createRequestBody(true);
+    let start = Date.now();
+    let response = await SELF.fetch(url, { method: 'POST', body: body, headers: { 'Content-Type': 'application/json', 'api-key': 'oilbeater' } });
+    const value = await response.text()
+    const duration = Date.now() - start
+
+    expect(response.status).toBe(200);
+    expect(response.headers.get('malacca-cache-status')).toBe('miss');
+
+    start = Date.now();
+    response = await SELF.fetch(url, {method: 'POST', body: body, headers: {'Content-Type': 'application/json', 'api-key': 'oilbeater'}});
+    const cacheValue = await response.text()
+    const cacheDuration = Date.now() - start
+
+    expect(response.status).toBe(200);
+    expect(response.headers.get('malacca-cache-status')).toBe('hit');
+    expect(value).toEqual(cacheValue)
+    expect(duration/2).toBeGreaterThan(cacheDuration)
+  });
+
 });