From 4d6c82deb91dd3dce83431012968f153442ba8fd Mon Sep 17 00:00:00 2001
From: "Q.A.zh" <40236765+QAbot-zh@users.noreply.github.com>
Date: Sat, 28 Dec 2024 09:09:41 +0000
Subject: [PATCH] add collectModelTable unit test

---
 test/model.test.ts | 143 +++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 143 insertions(+)
 create mode 100644 test/model.test.ts

diff --git a/test/model.test.ts b/test/model.test.ts
new file mode 100644
index 000000000..36f9e3472
--- /dev/null
+++ b/test/model.test.ts
@@ -0,0 +1,143 @@
+import { collectModelTable } from "@/app/utils/model"
+import { LLMModel,LLMModelProvider } from "@/app/client/api";
+
+describe('collectModelTable', () => {
+  const mockModels: readonly LLMModel[] = [
+    {
+      name: 'gpt-3.5-turbo',
+      available: true,
+      provider: {
+        id: 'openai',
+        providerName: 'OpenAI',
+        providerType: 'openai',
+      } as LLMModelProvider,
+      sorted: 1,
+    },
+    {
+      name: 'gpt-4',
+      available: true,
+      provider: {
+        id: 'openai',
+        providerName: 'OpenAI',
+        providerType: 'openai',
+      } as LLMModelProvider,
+      sorted: 1,
+    },
+    {
+      name: 'gpt-3.5-turbo',
+      available: true,
+      provider: {
+        id: 'azure',
+        providerName: 'Azure',
+        providerType: 'azure',
+      } as LLMModelProvider,
+      sorted: 2,
+    },
+    {
+      name: 'gpt-4',
+      available: true,
+      provider: {
+        id: 'azure',
+        providerName: 'Azure',
+        providerType: 'azure',
+      } as LLMModelProvider,
+      sorted: 2,
+    },
+    {
+      name: 'gemini-pro',
+      available: true,
+      provider: {
+        id: 'google',
+        providerName: 'Google',
+        providerType: 'google',
+      } as LLMModelProvider,
+      sorted: 3,
+    },
+    ,
+    {
+      name: 'claude-3-haiku-20240307',
+      available: true,
+      provider: {
+        id: 'anthropic',
+        providerName: 'Anthropic',
+        providerType: 'anthropic',
+      } as LLMModelProvider,
+      sorted: 4,
+    },
+    {
+      name: 'grok-beta',
+      available: true,
+      provider: {
+        id: 'xai',
+        providerName: 'XAI',
+        providerType: 'xai',
+      } as LLMModelProvider,
+      sorted: 11,
+    },
+  ];
+
+  test('all models shoule be available', () => {
+    const customModels = '';
+    const result = collectModelTable(mockModels, customModels);
+
+    expect(result['gpt-3.5@openai'].available).toBe(true);
+    expect(result['gpt-4@openai'].available).toBe(true);
+    expect(result['gpt-3.5@azure'].available).toBe(true);
+    expect(result['gpt-4@azure'].available).toBe(true);
+    expect(result['gemini-pro@google'].available).toBe(true);
+    expect(result['claude-3-haiku-20240307@claude'].available).toBe(true);
+    expect(result['grok-beta@xai'].available).toBe(true);
+  });
+  test('should exclude all models when custom is "-all"', () => {
+    const customModels = '-all';
+    const result = collectModelTable(mockModels, customModels);
+
+    expect(result['gpt-3.5@openai'].available).toBe(false);
+    expect(result['gpt-4@openai'].available).toBe(false);
+    expect(result['gpt-3.5@azure'].available).toBe(false);
+    expect(result['gpt-4@azure'].available).toBe(false);
+    expect(result['gemini-pro@google'].available).toBe(false);
+    expect(result['claude-3-haiku-20240307@anthropic'].available).toBe(false);
+    expect(result['grok-beta@xai'].available).toBe(false);
+  });
+
+  test('should exclude all Azure models when custom is "-*azure"', () => {
+    const customModels = '-*azure';
+    const result = collectModelTable(mockModels, customModels);
+
+    expect(result['gpt-3.5@openai'].available).toBe(true);
+    expect(result['gpt-4@openai'].available).toBe(true);
+    expect(result['gpt-3.5@azure'].available).toBe(false);
+    expect(result['gpt-4@azure'].available).toBe(false);
+    expect(result['gemini-pro@google'].available).toBe(true);
+    expect(result['claude-3-haiku-20240307@anthropic'].available).toBe(true);
+    expect(result['grok-beta@xai'].available).toBe(true);
+  });
+
+  test('should exclude Google and XAI models when custom is "-*google,-*xai"', () => {
+    const customModels = '-*google,-*xai';
+    const result = collectModelTable(mockModels, customModels);
+
+    expect(result['gpt-3.5@openai'].available).toBe(true);
+    expect(result['gpt-4@openai'].available).toBe(true);
+    expect(result['gpt-3.5@azure'].available).toBe(true);
+    expect(result['gpt-4@azure'].available).toBe(true);
+    expect(result['gemini-pro@google'].available).toBe(false);
+    expect(result['claude-3-haiku-20240307@anthropic'].available).toBe(true);
+    expect(result['grok-beta@xai'].available).toBe(false);
+  });
+
+  test('should exclude all models except OpenAI when custom is "-all, +*openai"', () => {
+    const customModels = '-all, +*openai,gpt-4o@azure';
+    const result = collectModelTable(mockModels, customModels);
+
+    expect(result['gpt-3.5@openai'].available).toBe(true);
+    expect(result['gpt-4@openai'].available).toBe(true);
+    expect(result['gpt-3.5@azure'].available).toBe(false);
+    expect(result['gpt-4@azure'].available).toBe(false);
+    expect(result['gemini-pro@google'].available).toBe(false);
+    expect(result['claude-3-haiku-20240307@anthropic'].available).toBe(false);
+    expect(result['grok-beta@xai'].available).toBe(false);
+    expect(result['gpt-4o@azure'].available).toBe(true);
+  });
+});
\ No newline at end of file