feat: Add comprehensive S3 storage testing with mocks
CRITICAL: Ensures production reliability at scale - Add aws-sdk-client-mock for S3 testing without AWS resources - Create comprehensive test suite for ALL S3 operations - Test statistics persistence, recovery, and partitioning - Test rate limiting, throttling, and retry logic - Test concurrent operations from multiple services - Test data integrity with checksums - Test batch operations and performance optimizations - Simulate millions of records for scale testing - Ensure backward compatibility with legacy formats Also: Remove deprecated addSmart() method (functionality merged into add()) These tests are critical for production deployments where S3 is the primary storage backend for scaled systems handling millions of records.
This commit is contained in:
parent
a292c4d3c9
commit
0ec21e96e6
5 changed files with 1700 additions and 8 deletions
448
tests/s3-statistics-critical.test.ts
Normal file
448
tests/s3-statistics-critical.test.ts
Normal file
|
|
@ -0,0 +1,448 @@
|
|||
/**
|
||||
* 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'
|
||||
|
||||
// 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({
|
||||
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({
|
||||
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({
|
||||
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({
|
||||
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({
|
||||
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({
|
||||
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({
|
||||
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({
|
||||
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({
|
||||
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)
|
||||
})
|
||||
})
|
||||
})
|
||||
Loading…
Add table
Add a link
Reference in a new issue