@boundless-oss/atlas/src/modules/rag-retrieval/__tests__/tools.test.ts

Atlas - MCP Server for comprehensive startup project management

import { describe, it, expect, beforeEach, afterEach } from 'vitest';
import { SQLiteManager } from '../../../storage/sqlite-manager.js';
import { setupRAGRetrievalTools } from '../tools.js';
import { RequestContext } from '../../../core/types.js';
import { randomUUID } from 'crypto';
import { promises as fs } from 'fs';
import path from 'path';
import os from 'os';

describe('RAG Retrieval Tools', () => {
  let db: SQLiteManager;
  let tools: any;
  let context: RequestContext;
  let tempDir: string;
  let testFile: string;

  beforeEach(async () => {
    // Create temporary database
    tempDir = await fs.mkdtemp(path.join(os.tmpdir(), 'rag-test-'));
    const dbPath = path.join(tempDir, 'test.db');
    
    db = new SQLiteManager(dbPath);
    await db.initialize();
    
    // Add a small delay to ensure database is fully ready
    await new Promise(resolve => setTimeout(resolve, 50));
    
    // Create a test project record to satisfy foreign key constraints
    const projectResult = await db.run(`
      INSERT OR REPLACE INTO projects (id, name, description, created_at, updated_at)
      VALUES (?, ?, ?, ?, ?)
    `, ['test-project', 'Test Project', 'Test project for RAG tests', Date.now(), Date.now()]);
    
    // Verify project was created
    if (!projectResult.success) {
      throw new Error('Failed to create test project: ' + projectResult.error);
    }
    
    // Setup context
    context = {
      requestId: randomUUID(),
      userId: 'test-user',
      projectId: 'test-project',
      db: db,
      startTime: Date.now()
    };
    
    // Load tools
    const registration = await setupRAGRetrievalTools();
    tools = registration.tools.reduce((acc: any, tool: any) => {
      acc[tool.name] = tool;
      return acc;
    }, {});
    
    // Create test markdown file
    testFile = path.join(tempDir, 'test-doc.md');
    await fs.writeFile(testFile, '# Test Document\n\nThis is a test document for RAG indexing.\n\n## Section 1\n\nSome content here.\n\n## Section 2\n\nMore content here.');
  });

  afterEach(async () => {
    await db.close();
    await fs.rm(tempDir, { recursive: true, force: true });
  });

  describe('rag_index_document', () => {
    it('should index a single document successfully', async () => {
      const result = await tools.rag_index_document.execute({
        path: path.relative(process.cwd(), testFile)
      }, context);

      expect(result.success).toBe(true);
      expect(result.data.document).toBeDefined();
      expect(result.data.document.path).toBeDefined();
      expect(result.data.document.chunkCount).toBeGreaterThan(0);
      expect(result.data.message).toContain('Successfully indexed document');
    });

    it('should handle non-existent file', async () => {
      const result = await tools.rag_index_document.execute({
        path: 'non-existent-file.md'
      }, context);

      expect(result.success).toBe(false);
      expect(result.error?.code).toBe('FILE_NOT_FOUND');
      expect(result.error?.message).toContain('Document not found');
    });

    it('should update existing document', async () => {
      // Index once
      const firstResult = await tools.rag_index_document.execute({
        path: path.relative(process.cwd(), testFile)
      }, context);
      expect(firstResult.success).toBe(true);

      // Update file content
      await fs.writeFile(testFile, '# Updated Document\n\nThis content has been updated.');

      // Index again
      const secondResult = await tools.rag_index_document.execute({
        path: path.relative(process.cwd(), testFile)
      }, context);
      
      expect(secondResult.success).toBe(true);
      expect(secondResult.data.document.id).toBeDefined();
    });
  });

  describe('rag_search', () => {
    beforeEach(async () => {
      // Add a small delay to ensure database is ready
      await new Promise(resolve => setTimeout(resolve, 10));
      
      // Index test document before search tests
      const indexResult = await tools.rag_index_document.execute({
        path: path.relative(process.cwd(), testFile)
      }, context);
      
      // Ensure indexing succeeded before running search tests
      expect(indexResult.success).toBe(true);
    });

    it('should search indexed documents', async () => {
      const result = await tools.rag_search.execute({
        query: 'test document',
        limit: 5
      }, context);

      expect(result.success).toBe(true);
      expect(result.data.results).toBeDefined();
      expect(Array.isArray(result.data.results)).toBe(true);
      expect(result.data.query).toBe('test document');
      expect(result.data.executionTime).toBeGreaterThanOrEqual(0);
    });

    it('should handle empty query gracefully', async () => {
      const result = await tools.rag_search.execute({
        query: 'nonexistent content that should not match'
      }, context);

      expect(result.success).toBe(true);
      expect(result.data.results).toBeDefined();
      expect(Array.isArray(result.data.results)).toBe(true);
    });

    it('should respect limit parameter', async () => {
      const result = await tools.rag_search.execute({
        query: 'content',
        limit: 2
      }, context);

      expect(result.success).toBe(true);
      expect(result.data.results.length).toBeLessThanOrEqual(2);
    });

    it('should log search history', async () => {
      await tools.rag_search.execute({
        query: 'test search history',
        limit: 5
      }, context);

      // Check that search was logged
      const historyResult = await db.query(
        'SELECT * FROM rag_search_history WHERE project_id = ? ORDER BY created_at DESC LIMIT 1',
        [context.projectId]
      );

      expect(historyResult.success).toBe(true);
      expect(historyResult.data?.length).toBe(1);
      expect(historyResult.data?.[0].query).toBe('test search history');
      expect(historyResult.data?.[0].limit_count).toBe(5);
    });
  });

  describe('rag_get_stats', () => {
    it('should return empty stats for new index', async () => {
      const result = await tools.rag_get_stats.execute({}, context);

      expect(result.success).toBe(true);
      expect(result.data.stats).toBeDefined();
      expect(result.data.stats.totalDocuments).toBe(0);
      expect(result.data.stats.totalChunks).toBe(0);
      expect(result.data.stats.totalCollections).toBe(0);
      expect(result.data.stats.lastIndexed).toBe('Never');
    });

    it('should return correct stats after indexing', async () => {
      // Index a document first
      await tools.rag_index_document.execute({
        path: path.relative(process.cwd(), testFile)
      }, context);

      const result = await tools.rag_get_stats.execute({}, context);

      expect(result.success).toBe(true);
      expect(result.data.stats).toBeDefined();
      expect(result.data.stats.totalDocuments).toBe(1);
      expect(result.data.stats.totalChunks).toBeGreaterThan(0);
      expect(result.data.stats.lastIndexed).not.toBe('Never');
    });
  });

  describe('rag_clear_index', () => {
    beforeEach(async () => {
      // Index test document before clear tests
      await tools.rag_index_document.execute({
        path: path.relative(process.cwd(), testFile)
      }, context);
    });

    it('should clear all indexed documents', async () => {
      // Verify there are documents
      let statsResult = await tools.rag_get_stats.execute({}, context);
      expect(statsResult.data.stats.totalDocuments).toBe(1);

      // Clear index
      const result = await tools.rag_clear_index.execute({}, context);
      expect(result.success).toBe(true);
      expect(result.data.message).toContain('Cleared all indexed documents');

      // Verify documents are cleared
      statsResult = await tools.rag_get_stats.execute({}, context);
      expect(statsResult.data.stats.totalDocuments).toBe(0);
    });
  });

  describe('rag_index_collection', () => {
    it('should create and index default docs collection', async () => {
      const result = await tools.rag_index_collection.execute({
        collection: 'docs'
      }, context);

      expect(result.success).toBe(true);
      expect(result.data.collection).toBeDefined();
      expect(result.data.collection.name).toBe('docs');
      expect(result.data.message).toContain('Indexed collection');
    });

    it('should create and index default readme collection', async () => {
      const result = await tools.rag_index_collection.execute({
        collection: 'readme'
      }, context);

      expect(result.success).toBe(true);
      expect(result.data.collection).toBeDefined();
      expect(result.data.collection.name).toBe('readme');
    });

    it('should handle non-existent collection', async () => {
      const result = await tools.rag_index_collection.execute({
        collection: 'non-existent-collection'
      }, context);

      expect(result.success).toBe(false);
      expect(result.error?.code).toBe('COLLECTION_NOT_FOUND');
      expect(result.error?.message).toContain('Collection not found');
    });
  });

  describe('rag_index_directory', () => {
    let testDir: string;

    beforeEach(async () => {
      testDir = path.join(tempDir, 'test-docs');
      await fs.mkdir(testDir, { recursive: true });
      
      // Create multiple test files
      await fs.writeFile(path.join(testDir, 'doc1.md'), '# Document 1\n\nContent of document 1.');
      await fs.writeFile(path.join(testDir, 'doc2.md'), '# Document 2\n\nContent of document 2.');
      await fs.writeFile(path.join(testDir, 'not-markdown.txt'), 'This should be ignored.');
    });

    it('should index all markdown files in directory', async () => {
      const result = await tools.rag_index_directory.execute({
        path: path.relative(process.cwd(), testDir)
      }, context);

      expect(result.success).toBe(true);
      expect(result.data.summary).toBeDefined();
      expect(result.data.summary.totalFiles).toBe(2); // Only .md files
      expect(result.data.summary.indexed).toBe(2);
      expect(result.data.summary.failed).toBe(0);
    });

    it('should handle non-existent directory', async () => {
      const result = await tools.rag_index_directory.execute({
        path: 'non-existent-directory'
      }, context);

      expect(result.success).toBe(false);
      expect(result.error?.code).toBe('DIRECTORY_NOT_FOUND');
    });

    it('should handle file instead of directory', async () => {
      const result = await tools.rag_index_directory.execute({
        path: path.relative(process.cwd(), testFile)
      }, context);

      expect(result.success).toBe(false);
      expect(result.error?.code).toBe('NOT_A_DIRECTORY');
    });
  });

  describe('Input validation', () => {
    it('should validate required parameters', async () => {
      // Test missing query for search
      expect(() => {
        tools.rag_search.inputSchema;
      }).not.toThrow();
      
      // Test missing path for index document
      expect(() => {
        tools.rag_index_document.inputSchema;
      }).not.toThrow();
    });

    it('should validate parameter types and constraints', async () => {
      const searchSchema = tools.rag_search.inputSchema;
      expect(searchSchema.properties.query.type).toBe('string');
      expect(searchSchema.properties.limit.minimum).toBe(1);
      expect(searchSchema.properties.limit.maximum).toBe(100);
      expect(searchSchema.properties.threshold.minimum).toBe(0);
      expect(searchSchema.properties.threshold.maximum).toBe(1);
    });
  });
});