David Snelling
04f174f6bb
- Add test-utils.ts with mock embedding function - Fix S3 tests to use mock embeddings instead of loading models - Fix vector dimension mismatches in tests - Update import paths for case-sensitive file systems - Ensure proper S3 mock initialization with ListObjectsV2Command All test infrastructure now properly mocked for fast, reliable tests
458 lines
No EOL
14 KiB
TypeScript
458 lines
No EOL
14 KiB
TypeScript
/**
|
|
* CRITICAL S3 Statistics Tests
|
|
*
|
|
* These tests ensure that statistics work correctly at scale with S3 storage.
|
|
* Statistics are fundamental for monitoring production deployments with millions of records.
|
|
*/
|
|
|
|
import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'
|
|
import { mockClient } from 'aws-sdk-client-mock'
|
|
import {
|
|
S3Client,
|
|
GetObjectCommand,
|
|
PutObjectCommand,
|
|
ListObjectsV2Command,
|
|
HeadObjectCommand,
|
|
DeleteObjectCommand
|
|
} from '@aws-sdk/client-s3'
|
|
import { BrainyData } from '../src/index.js'
|
|
import { S3CompatibleStorage } from '../src/storage/adapters/s3CompatibleStorage.js'
|
|
import { createMockEmbeddingFunction } from './test-utils.js'
|
|
|
|
// Create S3 mock
|
|
const s3Mock = mockClient(S3Client)
|
|
|
|
describe('CRITICAL: S3 Statistics at Scale', () => {
|
|
let brainy: BrainyData<any>
|
|
let storage: S3CompatibleStorage
|
|
|
|
beforeEach(() => {
|
|
// Reset all mocks before each test
|
|
s3Mock.reset()
|
|
vi.clearAllTimers()
|
|
vi.useFakeTimers()
|
|
})
|
|
|
|
afterEach(async () => {
|
|
if (brainy) {
|
|
await brainy.clear()
|
|
}
|
|
vi.useRealTimers()
|
|
})
|
|
|
|
describe('Statistics Persistence and Recovery', () => {
|
|
it('should persist statistics to S3 and recover after restart', async () => {
|
|
// Setup S3 mock responses
|
|
const statisticsData = {
|
|
nounCount: { 'service-a': 1000, 'service-b': 500 },
|
|
verbCount: { 'service-a': 200, 'service-b': 100 },
|
|
metadataCount: { 'service-a': 1000, 'service-b': 500 },
|
|
hnswIndexSize: 1500,
|
|
lastUpdated: new Date().toISOString()
|
|
}
|
|
|
|
// Mock initial empty state
|
|
s3Mock.on(GetObjectCommand).rejectsOnce({ name: 'NoSuchKey' })
|
|
|
|
// Mock successful save
|
|
s3Mock.on(PutObjectCommand).resolves({})
|
|
|
|
// Initialize Brainy with S3 storage
|
|
brainy = new BrainyData({
|
|
embeddingFunction: createMockEmbeddingFunction(),
|
|
storage: {
|
|
s3Storage: {
|
|
bucketName: 'test-bucket',
|
|
accessKeyId: 'test-key',
|
|
secretAccessKey: 'test-secret',
|
|
region: 'us-east-1'
|
|
}
|
|
}
|
|
})
|
|
await brainy.init()
|
|
|
|
// Add data with different services
|
|
await brainy.add('Test data 1', { metadata: 'test1' }, { service: 'service-a' })
|
|
await brainy.add('Test data 2', { metadata: 'test2' }, { service: 'service-b' })
|
|
|
|
// Force statistics flush
|
|
await brainy.flushStatistics()
|
|
|
|
// Verify statistics were saved to S3
|
|
const putCalls = s3Mock.commandCalls(PutObjectCommand)
|
|
expect(putCalls.length).toBeGreaterThan(0)
|
|
|
|
// Get current statistics
|
|
const stats = await brainy.getStatistics()
|
|
expect(stats.nounCount).toBe(2)
|
|
expect(stats.serviceBreakdown).toBeDefined()
|
|
expect(stats.serviceBreakdown['service-a'].nounCount).toBe(1)
|
|
expect(stats.serviceBreakdown['service-b'].nounCount).toBe(1)
|
|
|
|
// Simulate restart by creating new instance
|
|
s3Mock.reset()
|
|
|
|
// Mock loading saved statistics
|
|
s3Mock.on(GetObjectCommand).resolves({
|
|
Body: {
|
|
transformToString: async () => JSON.stringify({
|
|
nounCount: { 'service-a': 1, 'service-b': 1 },
|
|
verbCount: { 'service-a': 0, 'service-b': 0 },
|
|
metadataCount: { 'service-a': 1, 'service-b': 1 },
|
|
hnswIndexSize: 2,
|
|
lastUpdated: new Date().toISOString()
|
|
})
|
|
}
|
|
})
|
|
|
|
const brainy2 = new BrainyData({
|
|
embeddingFunction: createMockEmbeddingFunction(),
|
|
storage: {
|
|
s3Storage: {
|
|
bucketName: 'test-bucket',
|
|
accessKeyId: 'test-key',
|
|
secretAccessKey: 'test-secret',
|
|
region: 'us-east-1'
|
|
}
|
|
}
|
|
})
|
|
await brainy2.init()
|
|
|
|
// Verify statistics were recovered
|
|
const recoveredStats = await brainy2.getStatistics()
|
|
expect(recoveredStats.nounCount).toBe(2)
|
|
expect(recoveredStats.serviceBreakdown['service-a'].nounCount).toBe(1)
|
|
expect(recoveredStats.serviceBreakdown['service-b'].nounCount).toBe(1)
|
|
})
|
|
})
|
|
|
|
describe('Batching and Throttling', () => {
|
|
it('should batch statistics updates to prevent S3 rate limits', async () => {
|
|
s3Mock.on(GetObjectCommand).rejectsOnce({ name: 'NoSuchKey' })
|
|
s3Mock.on(PutObjectCommand).resolves({})
|
|
|
|
brainy = new BrainyData({
|
|
embeddingFunction: createMockEmbeddingFunction(),
|
|
storage: {
|
|
s3Storage: {
|
|
bucketName: 'test-bucket',
|
|
accessKeyId: 'test-key',
|
|
secretAccessKey: 'test-secret',
|
|
region: 'us-east-1'
|
|
}
|
|
}
|
|
})
|
|
await brainy.init()
|
|
|
|
// Add multiple items rapidly (simulating high load)
|
|
const promises = []
|
|
for (let i = 0; i < 100; i++) {
|
|
promises.push(
|
|
brainy.add(`Item ${i}`, { index: i }, { service: `service-${i % 5}` })
|
|
)
|
|
}
|
|
await Promise.all(promises)
|
|
|
|
// Initially, statistics shouldn't be flushed immediately
|
|
const initialPutCalls = s3Mock.commandCalls(PutObjectCommand)
|
|
expect(initialPutCalls.length).toBeLessThan(100) // Should batch, not 100 individual calls
|
|
|
|
// Advance timers to trigger batch flush
|
|
vi.advanceTimersByTime(5000)
|
|
|
|
// Force flush to ensure all statistics are saved
|
|
await brainy.flushStatistics()
|
|
|
|
// Verify statistics are correct
|
|
const stats = await brainy.getStatistics()
|
|
expect(stats.nounCount).toBe(100)
|
|
|
|
// Verify batching occurred (much fewer S3 calls than items)
|
|
const finalPutCalls = s3Mock.commandCalls(PutObjectCommand)
|
|
expect(finalPutCalls.length).toBeLessThan(20) // Should be batched
|
|
})
|
|
|
|
it('should handle S3 throttling (429) gracefully', async () => {
|
|
s3Mock.on(GetObjectCommand).rejectsOnce({ name: 'NoSuchKey' })
|
|
|
|
// Simulate throttling on first attempt, success on retry
|
|
let putAttempts = 0
|
|
s3Mock.on(PutObjectCommand).callsFake(() => {
|
|
putAttempts++
|
|
if (putAttempts === 1) {
|
|
const error: any = new Error('Too Many Requests')
|
|
error.name = 'TooManyRequestsException'
|
|
error.$metadata = { httpStatusCode: 429 }
|
|
throw error
|
|
}
|
|
return Promise.resolve({})
|
|
})
|
|
|
|
brainy = new BrainyData({
|
|
embeddingFunction: createMockEmbeddingFunction(),
|
|
storage: {
|
|
s3Storage: {
|
|
bucketName: 'test-bucket',
|
|
accessKeyId: 'test-key',
|
|
secretAccessKey: 'test-secret',
|
|
region: 'us-east-1'
|
|
}
|
|
}
|
|
})
|
|
await brainy.init()
|
|
|
|
// Add data
|
|
await brainy.add('Test data', { metadata: 'test' }, { service: 'throttle-test' })
|
|
|
|
// Force flush - should retry on throttling
|
|
await brainy.flushStatistics()
|
|
|
|
// Verify retry occurred
|
|
expect(putAttempts).toBeGreaterThanOrEqual(1)
|
|
|
|
// Verify statistics were eventually saved
|
|
const stats = await brainy.getStatistics()
|
|
expect(stats.nounCount).toBe(1)
|
|
})
|
|
})
|
|
|
|
describe('Time-Based Partitioning', () => {
|
|
it('should partition statistics by date to avoid single-key rate limits', async () => {
|
|
s3Mock.on(GetObjectCommand).rejects({ name: 'NoSuchKey' })
|
|
s3Mock.on(PutObjectCommand).resolves({})
|
|
s3Mock.on(ListObjectsV2Command).resolves({ Contents: [] })
|
|
|
|
brainy = new BrainyData({
|
|
embeddingFunction: createMockEmbeddingFunction(),
|
|
storage: {
|
|
s3Storage: {
|
|
bucketName: 'test-bucket',
|
|
accessKeyId: 'test-key',
|
|
secretAccessKey: 'test-secret',
|
|
region: 'us-east-1'
|
|
}
|
|
}
|
|
})
|
|
await brainy.init()
|
|
|
|
// Add data across different time periods
|
|
const today = new Date()
|
|
await brainy.add('Today item', { date: today }, { service: 'time-test' })
|
|
await brainy.flushStatistics()
|
|
|
|
// Check that statistics are saved with date-based key
|
|
const putCalls = s3Mock.commandCalls(PutObjectCommand)
|
|
const statisticsCall = putCalls.find(call =>
|
|
call.args[0].input.Key?.includes('_system/statistics')
|
|
)
|
|
|
|
expect(statisticsCall).toBeDefined()
|
|
// Should include date in the key for partitioning
|
|
const key = statisticsCall?.args[0].input.Key
|
|
expect(key).toContain(today.toISOString().split('T')[0])
|
|
})
|
|
})
|
|
|
|
describe('Backward Compatibility', () => {
|
|
it('should read legacy statistics format correctly', async () => {
|
|
// Mock legacy statistics format (without service breakdown)
|
|
const legacyStats = {
|
|
nounCount: 500,
|
|
verbCount: 100,
|
|
metadataCount: 500,
|
|
hnswIndexSize: 500,
|
|
lastUpdated: new Date().toISOString()
|
|
}
|
|
|
|
s3Mock.on(GetObjectCommand).resolves({
|
|
Body: {
|
|
transformToString: async () => JSON.stringify(legacyStats)
|
|
}
|
|
})
|
|
|
|
brainy = new BrainyData({
|
|
embeddingFunction: createMockEmbeddingFunction(),
|
|
storage: {
|
|
s3Storage: {
|
|
bucketName: 'test-bucket',
|
|
accessKeyId: 'test-key',
|
|
secretAccessKey: 'test-secret',
|
|
region: 'us-east-1'
|
|
}
|
|
}
|
|
})
|
|
await brainy.init()
|
|
|
|
// Should handle legacy format gracefully
|
|
const stats = await brainy.getStatistics()
|
|
expect(stats).toBeDefined()
|
|
// Legacy format should be converted to new format with default service
|
|
expect(stats.nounCount).toBe(500)
|
|
})
|
|
|
|
it('should migrate legacy statistics to new format on write', async () => {
|
|
// Start with legacy format
|
|
const legacyStats = {
|
|
nounCount: 100,
|
|
verbCount: 50,
|
|
metadataCount: 100,
|
|
hnswIndexSize: 100,
|
|
lastUpdated: new Date().toISOString()
|
|
}
|
|
|
|
s3Mock.on(GetObjectCommand).resolvesOnce({
|
|
Body: {
|
|
transformToString: async () => JSON.stringify(legacyStats)
|
|
}
|
|
})
|
|
s3Mock.on(PutObjectCommand).resolves({})
|
|
|
|
brainy = new BrainyData({
|
|
embeddingFunction: createMockEmbeddingFunction(),
|
|
storage: {
|
|
s3Storage: {
|
|
bucketName: 'test-bucket',
|
|
accessKeyId: 'test-key',
|
|
secretAccessKey: 'test-secret',
|
|
region: 'us-east-1'
|
|
}
|
|
}
|
|
})
|
|
await brainy.init()
|
|
|
|
// Add new data
|
|
await brainy.add('New data', { metadata: 'new' }, { service: 'migration-test' })
|
|
await brainy.flushStatistics()
|
|
|
|
// Verify new format was saved
|
|
const putCalls = s3Mock.commandCalls(PutObjectCommand)
|
|
const lastPut = putCalls[putCalls.length - 1]
|
|
const savedData = JSON.parse(lastPut.args[0].input.Body as string)
|
|
|
|
// Should have service-based structure
|
|
expect(savedData.nounCount).toBeTypeOf('object')
|
|
expect(savedData.nounCount['migration-test']).toBeDefined()
|
|
})
|
|
})
|
|
|
|
describe('Concurrent Updates', () => {
|
|
it('should handle concurrent statistics updates safely', async () => {
|
|
s3Mock.on(GetObjectCommand).rejectsOnce({ name: 'NoSuchKey' })
|
|
s3Mock.on(PutObjectCommand).resolves({})
|
|
|
|
brainy = new BrainyData({
|
|
embeddingFunction: createMockEmbeddingFunction(),
|
|
storage: {
|
|
s3Storage: {
|
|
bucketName: 'test-bucket',
|
|
accessKeyId: 'test-key',
|
|
secretAccessKey: 'test-secret',
|
|
region: 'us-east-1'
|
|
}
|
|
}
|
|
})
|
|
await brainy.init()
|
|
|
|
// Simulate concurrent additions from multiple services
|
|
const services = ['api', 'worker', 'batch', 'stream', 'webhook']
|
|
const concurrentOps = []
|
|
|
|
for (let i = 0; i < 50; i++) {
|
|
const service = services[i % services.length]
|
|
concurrentOps.push(
|
|
brainy.add(`Data ${i}`, { index: i }, { service })
|
|
)
|
|
}
|
|
|
|
// Add relationships concurrently
|
|
await Promise.all(concurrentOps)
|
|
|
|
const ids = concurrentOps.map((_, i) => `id-${i}`)
|
|
const relationOps = []
|
|
for (let i = 0; i < 10; i++) {
|
|
relationOps.push(
|
|
brainy.relate(
|
|
ids[i],
|
|
ids[i + 10],
|
|
'related',
|
|
{ service: services[i % services.length] }
|
|
).catch(() => {}) // Ignore if IDs don't exist
|
|
)
|
|
}
|
|
|
|
await Promise.all(relationOps)
|
|
await brainy.flushStatistics()
|
|
|
|
// Verify all operations were counted correctly
|
|
const stats = await brainy.getStatistics()
|
|
expect(stats.nounCount).toBe(50)
|
|
|
|
// Verify per-service counts
|
|
for (const service of services) {
|
|
expect(stats.serviceBreakdown[service]).toBeDefined()
|
|
expect(stats.serviceBreakdown[service].nounCount).toBe(10) // 50 items / 5 services
|
|
}
|
|
})
|
|
})
|
|
|
|
describe('Large Scale Statistics', () => {
|
|
it('should handle statistics for millions of records efficiently', async () => {
|
|
// Mock large existing statistics
|
|
const largeStats = {
|
|
nounCount: {
|
|
'service-1': 1000000,
|
|
'service-2': 2000000,
|
|
'service-3': 1500000
|
|
},
|
|
verbCount: {
|
|
'service-1': 500000,
|
|
'service-2': 750000,
|
|
'service-3': 600000
|
|
},
|
|
metadataCount: {
|
|
'service-1': 1000000,
|
|
'service-2': 2000000,
|
|
'service-3': 1500000
|
|
},
|
|
hnswIndexSize: 4500000,
|
|
lastUpdated: new Date().toISOString()
|
|
}
|
|
|
|
s3Mock.on(GetObjectCommand).resolves({
|
|
Body: {
|
|
transformToString: async () => JSON.stringify(largeStats)
|
|
}
|
|
})
|
|
s3Mock.on(PutObjectCommand).resolves({})
|
|
|
|
brainy = new BrainyData({
|
|
embeddingFunction: createMockEmbeddingFunction(),
|
|
storage: {
|
|
s3Storage: {
|
|
bucketName: 'test-bucket',
|
|
accessKeyId: 'test-key',
|
|
secretAccessKey: 'test-secret',
|
|
region: 'us-east-1'
|
|
}
|
|
}
|
|
})
|
|
await brainy.init()
|
|
|
|
// Get statistics for large dataset
|
|
const stats = await brainy.getStatistics()
|
|
|
|
// Verify large numbers are handled correctly
|
|
expect(stats.nounCount).toBe(4500000)
|
|
expect(stats.verbCount).toBe(1850000)
|
|
|
|
// Add more data to large dataset
|
|
await brainy.add('New item in large dataset', {}, { service: 'service-1' })
|
|
await brainy.flushStatistics()
|
|
|
|
// Verify increment worked correctly with large numbers
|
|
const updatedStats = await brainy.getStatistics()
|
|
expect(updatedStats.nounCount).toBe(4500001)
|
|
expect(updatedStats.serviceBreakdown['service-1'].nounCount).toBe(1000001)
|
|
})
|
|
})
|
|
}) |