at-mintel/packages/memory-mcp/src/qdrant.test.ts

import { describe, it, expect, vi, beforeEach } from 'vitest';
import { QdrantMemoryService } from './qdrant.js';

vi.mock('@xenova/transformers', () => {
    return {
        env: { allowRemoteModels: false, localModelPath: './models' },
        pipeline: vi.fn().mockResolvedValue(async (text: string) => {
            // Mock embedding generation: returns an array of 384 numbers
            return { data: new Float32Array(384).fill(0.1) };
        }),
    };
});

const mockCreateCollection = vi.fn();
const mockGetCollections = vi.fn().mockResolvedValue({ collections: [] });
const mockUpsert = vi.fn();
const mockSearch = vi.fn().mockResolvedValue([
    {
        id: 'test-id',
        version: 1,
        score: 0.9,
        payload: { label: 'Test Label', content: 'Test Content' }
    }
]);

vi.mock('@qdrant/js-client-rest', () => {
    return {
        QdrantClient: vi.fn().mockImplementation(() => {
            return {
                getCollections: mockGetCollections,
                createCollection: mockCreateCollection,
                upsert: mockUpsert,
                search: mockSearch
            };
        })
    };
});

describe('QdrantMemoryService', () => {
    let service: QdrantMemoryService;

    beforeEach(() => {
        vi.clearAllMocks();
        service = new QdrantMemoryService('http://localhost:6333');
    });

    it('should initialize and create collection if missing', async () => {
        mockGetCollections.mockResolvedValueOnce({ collections: [] });
        await service.initialize();

        expect(mockGetCollections).toHaveBeenCalled();
        expect(mockCreateCollection).toHaveBeenCalledWith('mcp_memory', expect.any(Object));
    });

    it('should not create collection if it already exists', async () => {
        mockGetCollections.mockResolvedValueOnce({ collections: [{ name: 'mcp_memory' }] });
        await service.initialize();

        expect(mockCreateCollection).not.toHaveBeenCalled();
    });

    it('should store memory', async () => {
        await service.initialize();
        const result = await service.storeMemory('Design', 'Composition over Inheritance');

        expect(result).toBe(true);
        expect(mockUpsert).toHaveBeenCalledWith('mcp_memory', expect.objectContaining({
            wait: true,
            points: expect.arrayContaining([
                expect.objectContaining({
                    payload: expect.objectContaining({
                        label: 'Design',
                        content: 'Composition over Inheritance'
                    })
                })
            ])
        }));
    });

    it('should retrieve memory', async () => {
        await service.initialize();
        const results = await service.retrieveMemory('Design');

        expect(results).toHaveLength(1);
        expect(results[0].label).toBe('Test Label');
        expect(results[0].content).toBe('Test Content');
        expect(results[0].score).toBe(0.9);
    });
});