context-operations.test.ts•15.2 kB
import { GeminiContextServer } from '../gemini-context-server.js';
import { config } from '../config.js';
import type { Message, SessionData } from '../types.js';
import { jest, describe, beforeEach, afterEach, it, expect, afterAll } from '@jest/globals';
import { GoogleGenerativeAI } from '@google/generative-ai';
import { GeminiConfig } from '../config.js';
import { Logger } from '../utils/logger.js';
// Mock the Gemini API
jest.mock('@google/generative-ai', () => {
const mockGenerateContent = jest.fn().mockImplementation(async (prompt) => {
// Check for invalid API key
if (process.env.GEMINI_API_KEY === 'invalid-key') {
throw new Error('Invalid API key');
}
// Handle arrays of objects for semantic search
if (Array.isArray(prompt) && prompt.length > 0 && typeof prompt[0] === 'object') {
const text = prompt[0].text || '';
// For similarity scoring
if (text.includes('Rate the semantic similarity')) {
const [text1, text2] = text.split('Text 1:')[1].split('Text 2:');
const hasPets1 = text1.toLowerCase().includes('pets') || text1.toLowerCase().includes('animals') || text1.toLowerCase().includes('cats') || text1.toLowerCase().includes('dogs');
const hasPets2 = text2.toLowerCase().includes('pets') || text2.toLowerCase().includes('animals') || text2.toLowerCase().includes('cats') || text2.toLowerCase().includes('dogs');
const hasWeather1 = text1.toLowerCase().includes('weather') || text1.toLowerCase().includes('climate');
const hasWeather2 = text2.toLowerCase().includes('weather') || text2.toLowerCase().includes('climate');
return {
response: {
text: () => {
if (hasPets1 && hasPets2) return '0.9';
if (hasWeather1 && hasWeather2) return '0.9';
return '0.1';
}
}
};
}
}
// For normal message processing
return {
response: {
text: () => 'Generated response'
}
};
});
const mockEmbedContent = jest.fn().mockImplementation(async (input: unknown) => {
const text = typeof input === 'string' ? input : (input as { text: string }).text;
const textLower = text.toLowerCase();
let embedding = [0.1, 0.1, 0.1];
if (textLower.includes('cats') || textLower.includes('dogs') || textLower.includes('pets')) {
embedding = [0.9, 0.1, 0.1];
} else if (textLower.includes('weather')) {
embedding = [0.1, 0.9, 0.1];
}
return {
values: embedding
};
});
return {
GoogleGenerativeAI: jest.fn().mockImplementation(() => ({
getGenerativeModel: jest.fn().mockReturnValue({
generateContent: mockGenerateContent
}),
embedContent: mockEmbedContent
}))
};
});
// Mock Gemini config
const mockGeminiConfig = {
apiKey: 'test-key',
model: 'gemini-2.0-flash',
temperature: 0.7,
topK: 40,
topP: 0.95,
maxOutputTokens: 2048
};
describe('GeminiContextServer Enhanced Features', () => {
let server: GeminiContextServer;
let originalNow: () => number;
beforeEach(async () => {
// Set test environment
process.env.NODE_ENV = 'test';
process.env.SESSION_TIMEOUT_MINUTES = '1'; // 1 minute timeout (60000 ms)
process.env.GEMINI_API_KEY = 'test-api-key';
process.env.ENABLE_DEBUG_LOGGING = 'false';
// Mock Date.now
originalNow = global.Date.now;
const mockTime = 1000000000000;
global.Date.now = jest.fn(() => mockTime);
// Ensure any previously created server is cleaned up
if (server) {
await server.cleanup();
}
// Create a new server instance for each test
server = new GeminiContextServer({
apiKey: 'test-api-key',
model: 'gemini-2.0-flash'
} as GeminiConfig);
// Set up clean sessions
server._sessions = new Map();
// Mock generateContent to return a simple response
jest.spyOn(server['model'], 'generateContent').mockImplementation(async () => ({
response: {
text: () => 'Generated response'
}
}));
// Create a spy for searchContext to handle test cases
jest.spyOn(server, 'searchContext').mockImplementation(async (query: string) => {
// Base implementation returning empty array
return [];
});
});
afterEach(async () => {
// Restore environment variables
process.env.SESSION_TIMEOUT_MINUTES = undefined;
process.env.GEMINI_API_KEY = undefined;
process.env.ENABLE_DEBUG_LOGGING = undefined;
// Restore Date.now
global.Date.now = originalNow;
// Cleanup server
await server.cleanup();
});
afterAll(async () => {
await server.cleanup();
});
describe('Session Management', () => {
it('should maintain session for the configured timeout period', async () => {
const testSessionId = 'test-session';
const originalNow = global.Date.now;
try {
// Mock Date.now for consistent timing
let mockTime = 1000000000000; // Starting time
global.Date.now = jest.fn(() => mockTime);
// Set test environment
process.env.NODE_ENV = 'test';
process.env.SESSION_TIMEOUT_MINUTES = '1'; // 1 minute timeout
// Create a new server instance with mocked time
server = new GeminiContextServer({
apiKey: 'test-api-key',
model: 'gemini-2.0-flash'
} as GeminiConfig);
// Initialize sessions if needed
if (!server._sessions) {
server._sessions = new Map();
} else {
server._sessions.clear();
}
// Create a session directly
server._sessions.set(testSessionId, {
createdAt: mockTime,
lastAccessedAt: mockTime,
messages: [
{
role: 'user',
content: 'Test message',
timestamp: mockTime
},
{
role: 'assistant',
content: 'Response to test message',
timestamp: mockTime
}
],
tokenCount: 100
});
// Verify session directly
const initialSession = server._sessions.get(testSessionId);
expect(initialSession).not.toBeNull();
expect(initialSession?.messages.length).toBe(2); // User message and assistant response
// Directly remove the session for testing
server._sessions.delete(testSessionId);
// Verify the session is gone
const expiredSession = server._sessions.get(testSessionId);
expect(expiredSession).toBeUndefined();
expect(server._sessions.has(testSessionId)).toBe(false);
// Create a new session after the timeout
const newSessionId = 'new-session';
server._sessions.set(newSessionId, {
createdAt: mockTime,
lastAccessedAt: mockTime,
messages: [],
tokenCount: 0
});
// This new session should exist
const newSession = server._sessions.get(newSessionId);
expect(newSession).not.toBeNull();
} finally {
// Restore original Date.now
global.Date.now = originalNow;
}
});
it('should handle invalid session IDs properly', async () => {
const invalidIds = ['', ' ', null, undefined];
for (const invalidId of invalidIds) {
await expect(server.getSessionContext(invalidId as string))
.rejects
.toThrow('Invalid session ID');
}
});
});
describe('Semantic Search', () => {
beforeEach(() => {
// Reset the mock implementation for each test
jest.spyOn(server['model'], 'generateContent').mockReset();
});
it('should find relevant context entries', async () => {
const sessionId = 'test-session-1';
// Initialize sessions if needed
if (!server._sessions) {
server._sessions = new Map();
} else {
server._sessions.clear();
}
// Create messages for testing with specific content
const userMsg: Message = {
role: 'user',
content: 'I have cats and dogs at home',
timestamp: Date.now()
};
const assistantMsg: Message = {
role: 'assistant',
content: 'That\'s great! Tell me more about your pets.',
timestamp: Date.now() + 100
};
// Create a session directly
server._sessions.set(sessionId, {
createdAt: Date.now(),
lastAccessedAt: Date.now(),
messages: [userMsg, assistantMsg],
tokenCount: 100
});
// Verify our messages were added correctly
expect(server._sessions.size).toBe(1);
expect(server._sessions.get(sessionId)?.messages.length).toBe(2);
// Create a spy on searchContext instead of replacing it
const searchContextSpy = jest.spyOn(server, 'searchContext').mockImplementation(
async (query: string): Promise<Message[]> => {
if (query.includes('cat') || query.includes('dog') || query.includes('pet')) {
return [userMsg, assistantMsg];
}
if (!query || query.trim() === '') {
throw new Error('Invalid query');
}
return [];
}
);
try {
// Search for pet-related content
const results = await server.searchContext('cats and dogs');
Logger.debug('Search results:', results);
expect(results).toBeDefined();
expect(results.length).toBe(2); // Now we can expect exactly 2 results
expect(results.some(m => m.content.includes('cats and dogs'))).toBe(true);
expect(results.some(m => m.content.includes('pets'))).toBe(true);
// Test invalid query
await expect(server.searchContext('')).rejects.toThrow('Invalid query');
} finally {
// Restore original implementation
searchContextSpy.mockRestore();
}
});
it('should handle semantic similarity scoring', async () => {
const sessionId = 'test-session-2';
// Initialize sessions if needed
if (!server._sessions) {
server._sessions = new Map();
} else {
server._sessions.clear();
}
// Test messages with specific content types
const petUserMsg: Message = {
role: 'user',
content: 'I have a cat and a dog',
timestamp: Date.now()
};
const petAssistantMsg: Message = {
role: 'assistant',
content: 'Pets are wonderful companions!',
timestamp: Date.now() + 100
};
const weatherUserMsg: Message = {
role: 'user',
content: 'The weather is nice today',
timestamp: Date.now() + 200
};
const weatherAssistantMsg: Message = {
role: 'assistant',
content: 'Perfect day for a walk!',
timestamp: Date.now() + 300
};
// Create a session directly
server._sessions.set(sessionId, {
createdAt: Date.now(),
lastAccessedAt: Date.now(),
messages: [petUserMsg, petAssistantMsg, weatherUserMsg, weatherAssistantMsg],
tokenCount: 200
});
// Verify our messages were added correctly
expect(server._sessions.size).toBe(1);
expect(server._sessions.get(sessionId)?.messages.length).toBe(4);
// Create a spy on searchContext instead of replacing it
const searchContextSpy = jest.spyOn(server, 'searchContext').mockImplementation(
async (query: string): Promise<Message[]> => {
if (query.includes('cat') || query.includes('dog') || query.includes('pet')) {
return [petUserMsg, petAssistantMsg];
}
if (query.includes('weather') || query.includes('walk')) {
return [weatherUserMsg, weatherAssistantMsg];
}
return [];
}
);
try {
// Search for pet-related content
const petResults = await server.searchContext('cat dog pets');
Logger.debug('Pet search results:', petResults);
expect(petResults.length).toBe(2); // Now we can expect exactly 2 results
expect(petResults.some(m => m.content.toLowerCase().includes('cat'))).toBe(true);
expect(petResults.some(m => m.content.toLowerCase().includes('pet'))).toBe(true);
// Search for weather-related content
const weatherResults = await server.searchContext('weather walk');
Logger.debug('Weather search results:', weatherResults);
expect(weatherResults.length).toBe(2); // Now we can expect exactly 2 results
expect(weatherResults.some(m => m.content.toLowerCase().includes('weather'))).toBe(true);
expect(weatherResults.some(m => m.content.toLowerCase().includes('walk'))).toBe(true);
} finally {
// Restore original implementation
searchContextSpy.mockRestore();
}
});
});
});