From fb65ac42ea4b2aed0e23608e4faafbb06d8f6ca0 Mon Sep 17 00:00:00 2001 From: dal Date: Fri, 8 Aug 2025 09:32:53 -0600 Subject: [PATCH] test: update integration tests to include additional properties for agent options - Added userId, chatId, organizationId, dataSourceId, and dataSourceSyntax to the integration tests for the Think and Prep Agent and Mark Message Complete Step. - Cleaned up import statements for better organization and readability. - Ensured that the tests reflect the latest schema changes for improved tracking and functionality. These updates enhance the clarity and robustness of the integration tests, aligning them with recent changes in agent options. --- .../think-and-prep-agent.int.test.ts | 4 +++ .../mark-message-complete-step.int.test.ts | 35 +++++++------------ .../think-and-prep-step.test.ts | 20 +++++++++++ 3 files changed, 37 insertions(+), 22 deletions(-) diff --git a/packages/ai/src/agents/think-and-prep-agent/think-and-prep-agent.int.test.ts b/packages/ai/src/agents/think-and-prep-agent/think-and-prep-agent.int.test.ts index a77d7be3e..dda75f9a9 100644 --- a/packages/ai/src/agents/think-and-prep-agent/think-and-prep-agent.int.test.ts +++ b/packages/ai/src/agents/think-and-prep-agent/think-and-prep-agent.int.test.ts @@ -13,6 +13,10 @@ describe('Think and Prep Agent Integration Tests', () => { messageId: 'test-message-123', sql_dialect_guidance: 'postgresql', chatId: 'test-chat-123', + organizationId: 'test-organization-123', + userId: 'test-user-123', + dataSourceId: 'test-data-source-123', + dataSourceSyntax: 'postgresql', }); const streamResult = await thinkAndPrepAgent.stream({ diff --git a/packages/ai/src/steps/analyst-agent-steps/mark-message-complete-step/mark-message-complete-step.int.test.ts b/packages/ai/src/steps/analyst-agent-steps/mark-message-complete-step/mark-message-complete-step.int.test.ts index 35b5e3586..329a548d3 100644 --- a/packages/ai/src/steps/analyst-agent-steps/mark-message-complete-step/mark-message-complete-step.int.test.ts +++ b/packages/ai/src/steps/analyst-agent-steps/mark-message-complete-step/mark-message-complete-step.int.test.ts @@ -1,7 +1,7 @@ -import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; -import { db, messages, chats } from '@buster/database'; -import { eq } from 'drizzle-orm'; +import { chats, db, messages } from '@buster/database'; import { createTestChat, createTestMessage } from '@buster/test-utils'; +import { eq } from 'drizzle-orm'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { markMessageComplete } from './mark-message-complete-step'; describe('markMessageComplete integration test', () => { @@ -20,7 +20,7 @@ describe('markMessageComplete integration test', () => { // Create a test message - returns just the ID string testMessageId = await createTestMessage(testChatId, testUserId, { isCompleted: false, // Start with uncompleted message for testing - finalReasoningMessage: null, + finalReasoningMessage: '', }); }); @@ -40,11 +40,8 @@ describe('markMessageComplete integration test', () => { expect(result.completedAt).toBeDefined(); // Verify the message was updated in the database - const [updatedMessage] = await db - .select() - .from(messages) - .where(eq(messages.id, testMessageId)); - + const [updatedMessage] = await db.select().from(messages).where(eq(messages.id, testMessageId)); + expect(updatedMessage?.isCompleted).toBe(true); expect(updatedMessage?.finalReasoningMessage).toBe('Task completed successfully'); }); @@ -67,14 +64,11 @@ describe('markMessageComplete integration test', () => { expect(result.messageId).toBe(testMessageId); // Wait a moment for database update to complete - await new Promise(resolve => setTimeout(resolve, 100)); + await new Promise((resolve) => setTimeout(resolve, 100)); // Verify the chat was updated with file information - const [updatedChat] = await db - .select() - .from(chats) - .where(eq(chats.id, testChatId)); - + const [updatedChat] = await db.select().from(chats).where(eq(chats.id, testChatId)); + expect(updatedChat).toBeDefined(); expect(updatedChat?.mostRecentFileId).toBe(selectedFile.fileId); expect(updatedChat?.mostRecentFileType).toBe(selectedFile.fileType); @@ -106,7 +100,7 @@ describe('markMessageComplete integration test', () => { .select() .from(messages) .where(eq(messages.id, testMessageId)); - + expect(updatedMessage?.finalReasoningMessage).toBe('complete'); } }); @@ -128,11 +122,8 @@ describe('markMessageComplete integration test', () => { expect(result.success).toBe(true); // Verify the chat was not updated - const [chat] = await db - .select() - .from(chats) - .where(eq(chats.id, testChatId)); - + const [chat] = await db.select().from(chats).where(eq(chats.id, testChatId)); + expect(chat?.mostRecentFileId).toBeNull(); }); -}); \ No newline at end of file +}); diff --git a/packages/ai/src/steps/analyst-agent-steps/think-and-prep-step/think-and-prep-step.test.ts b/packages/ai/src/steps/analyst-agent-steps/think-and-prep-step/think-and-prep-step.test.ts index 2d5a4adba..564440993 100644 --- a/packages/ai/src/steps/analyst-agent-steps/think-and-prep-step/think-and-prep-step.test.ts +++ b/packages/ai/src/steps/analyst-agent-steps/think-and-prep-step/think-and-prep-step.test.ts @@ -39,6 +39,11 @@ describe('runThinkAndPrepAgentStep', () => { options: { messageId: 'test-message-id', sql_dialect_guidance: 'postgres', + userId: 'test-user-id', + chatId: 'test-chat-id', + organizationId: 'test-organization-id', + dataSourceId: 'test-data-source-id', + dataSourceSyntax: 'test-data-source-syntax', }, streamOptions: { messages: [{ role: 'user', content: 'Test prompt' }], @@ -71,6 +76,11 @@ describe('runThinkAndPrepAgentStep', () => { options: { messageId: 'test-message-id', sql_dialect_guidance: 'postgres', + userId: 'test-user-id', + chatId: 'test-chat-id', + organizationId: 'test-organization-id', + dataSourceId: 'test-data-source-id', + dataSourceSyntax: 'test-data-source-syntax', }, streamOptions: { messages: [{ role: 'user', content: 'Test prompt' }], @@ -100,6 +110,11 @@ describe('runThinkAndPrepAgentStep', () => { options: { messageId: 'test-message-id', sql_dialect_guidance: 'postgres', + userId: 'test-user-id', + chatId: 'test-chat-id', + organizationId: 'test-organization-id', + dataSourceId: 'test-data-source-id', + dataSourceSyntax: 'test-data-source-syntax', }, streamOptions: { messages: [{ role: 'user', content: 'Test prompt' }], @@ -127,6 +142,11 @@ describe('runThinkAndPrepAgentStep', () => { options: { messageId: 'test-message-id', sql_dialect_guidance: 'postgres', + userId: 'test-user-id', + chatId: 'test-chat-id', + organizationId: 'test-organization-id', + dataSourceId: 'test-data-source-id', + dataSourceSyntax: 'test-data-source-syntax', }, streamOptions: { messages: [{ role: 'user', content: 'Test prompt' }],