feat(chore): fix tests

Author: akshatdubeysf

src/__tests__/db-query/acceptance/dataset.controller.acceptance.ts

@@ -6,7 +6,7 @@ import {DataSet} from '../../../components/db-query/models';
 import {DataSetRepository} from '../../../components/db-query/repositories';
 import {AiIntegrationBindings} from '../../../keys';
 import {PermissionKey} from '../../../permissions';
-import {EmbeddingProvider, LLMProvider} from '../../../types';
+import {EmbeddingProvider} from '../../../types';
 import {TestApp} from '../../fixtures/test-app';
 import {buildToken, seedEmployees, setupApplication} from '../../test-helper';
 
@@ -19,22 +19,14 @@ describe('DatasetController', () => {
   before('setupApplication', async () => {
     ({app, client} = await setupApplication({
       noKnowledgeGraph: true,
+      llmStub: sinon.stub(),
     }));
     llmStub = sinon.stub();
     llmStub.resolves([[0.1, 0.2, 0.3]]);
 
     app.bind(AiIntegrationBindings.EmbeddingModel).to({
       embedDocuments: llmStub,
     } as unknown as EmbeddingProvider);
-    app
-      .bind(AiIntegrationBindings.SmartLLM)
-      .to(llmStub as unknown as LLMProvider);
-    app
-      .bind(AiIntegrationBindings.CheapLLM)
-      .to(llmStub as unknown as LLMProvider);
-    app
-      .bind(AiIntegrationBindings.FileLLM)
-      .to(llmStub as unknown as LLMProvider);
     await seedEmployees(app);
     await seedDataset(app);
   });

src/__tests__/fixtures/test-app.ts

@@ -44,12 +44,22 @@ export class TestApp extends BootMixin(
       this.bind(AiIntegrationBindings.CheapLLM).toProvider(Ollama);
       this.bind(AiIntegrationBindings.SmartLLM).toProvider(Ollama);
       this.bind(AiIntegrationBindings.FileLLM).toProvider(Ollama);
-    } else {
+      this.bind(AiIntegrationBindings.EmbeddingModel).toProvider(
+        OllamaEmbedding,
+      );
+    } else if (process.env.CEREBRAS === '1') {
       this.bind(AiIntegrationBindings.CheapLLM).toProvider(Cerebras);
       this.bind(AiIntegrationBindings.SmartLLM).toProvider(Cerebras);
       this.bind(AiIntegrationBindings.FileLLM).toProvider(Cerebras);
+      this.bind(AiIntegrationBindings.EmbeddingModel).toProvider(
+        OllamaEmbedding,
+      );
+    } else if (options.llmStub) {
+      this.bind(AiIntegrationBindings.CheapLLM).to(options.llmStub);
+      this.bind(AiIntegrationBindings.SmartLLM).to(options.llmStub);
+      this.bind(AiIntegrationBindings.FileLLM).to(options.llmStub);
+      this.bind(AiIntegrationBindings.EmbeddingModel).to(options.llmStub);
     }
-    this.bind(AiIntegrationBindings.EmbeddingModel).toProvider(OllamaEmbedding);
     this.bind('datasources.db').to(
       new juggler.DataSource({
         connector: 'sqlite3',

src/__tests__/test-helper.ts

@@ -27,6 +27,7 @@ config();
 
 export async function setupApplication(options: {
   noKnowledgeGraph?: boolean;
+  llmStub?: sinon.SinonStub;
 }): Promise<AppWithClient> {
   const restConfig = givenHttpServerConfig({
     // Customize the server configuration here.