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
654 lines
No EOL
19 KiB
TypeScript
654 lines
No EOL
19 KiB
TypeScript
/**
|
|
* COMPREHENSIVE S3 Storage Tests
|
|
*
|
|
* This test suite covers ALL S3-based features to ensure production reliability at scale.
|
|
* Tests include: statistics, nouns, verbs, metadata, HNSW index, caching, and error handling.
|
|
*/
|
|
|
|
import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest'
|
|
import { mockClient } from 'aws-sdk-client-mock'
|
|
import {
|
|
S3Client,
|
|
GetObjectCommand,
|
|
PutObjectCommand,
|
|
ListObjectsV2Command,
|
|
HeadObjectCommand,
|
|
DeleteObjectCommand,
|
|
DeleteObjectsCommand
|
|
} from '@aws-sdk/client-s3'
|
|
import { BrainyData } from '../src/index.js'
|
|
import { createMockEmbeddingFunction, createMockS3Body } from './test-utils.js'
|
|
|
|
// Create S3 mock
|
|
const s3Mock = mockClient(S3Client)
|
|
|
|
// Use the shared mock body helper
|
|
const createMockBody = createMockS3Body
|
|
|
|
describe('COMPREHENSIVE S3 Storage Tests', () => {
|
|
let brainy: BrainyData<any>
|
|
|
|
beforeEach(() => {
|
|
// Reset all mocks before each test
|
|
s3Mock.reset()
|
|
vi.clearAllTimers()
|
|
vi.useFakeTimers()
|
|
})
|
|
|
|
afterEach(async () => {
|
|
if (brainy) {
|
|
await brainy.clear()
|
|
}
|
|
vi.useRealTimers()
|
|
})
|
|
|
|
describe('Core S3 Operations', () => {
|
|
describe('Nouns (Vector Data)', () => {
|
|
it('should save and retrieve nouns from S3', async () => {
|
|
const nounData = {
|
|
id: 'test-noun-1',
|
|
vector: [0.1, 0.2, 0.3],
|
|
connections: {},
|
|
level: 0
|
|
}
|
|
|
|
// Mock S3 responses
|
|
s3Mock.on(GetObjectCommand, {
|
|
Key: 'nouns/test-noun-1.json'
|
|
}).resolves({
|
|
Body: createMockBody(nounData)
|
|
})
|
|
|
|
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 a noun (use string so it gets embedded)
|
|
const id = await brainy.add('Test noun content', { name: 'Test noun' })
|
|
|
|
// Verify noun was saved to S3
|
|
const putCalls = s3Mock.commandCalls(PutObjectCommand)
|
|
const nounSave = putCalls.find(call =>
|
|
call.args[0].input.Key?.includes('nouns/')
|
|
)
|
|
expect(nounSave).toBeDefined()
|
|
|
|
// Retrieve the noun
|
|
const retrieved = await brainy.get(id)
|
|
expect(retrieved).toBeDefined()
|
|
expect(retrieved?.metadata?.name).toBe('Test noun')
|
|
})
|
|
|
|
it('should handle batch noun operations efficiently', 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 batch of nouns
|
|
const items = Array(100).fill(0).map((_, i) => ({
|
|
data: `Item ${i}`,
|
|
metadata: { index: i }
|
|
}))
|
|
|
|
const ids = await brainy.addBatch(
|
|
items.map(item => item.data),
|
|
items.map(item => item.metadata)
|
|
)
|
|
|
|
expect(ids.length).toBe(100)
|
|
|
|
// Verify batch operations are optimized
|
|
const putCalls = s3Mock.commandCalls(PutObjectCommand)
|
|
// Should batch operations, not 100 individual calls
|
|
expect(putCalls.length).toBeLessThan(200) // Some batching should occur
|
|
})
|
|
})
|
|
|
|
describe('Verbs (Relationships)', () => {
|
|
it('should save and retrieve verbs from S3', async () => {
|
|
const verbData = {
|
|
id: 'verb-1',
|
|
source: 'noun-1',
|
|
target: 'noun-2',
|
|
type: 'relates_to',
|
|
vector: [0.4, 0.5, 0.6],
|
|
weight: 0.8
|
|
}
|
|
|
|
s3Mock.on(GetObjectCommand).rejects({ name: 'NoSuchKey' })
|
|
s3Mock.on(PutObjectCommand).resolves({})
|
|
s3Mock.on(ListObjectsV2Command).resolves({
|
|
Contents: [{
|
|
Key: 'verbs/verb-1.json'
|
|
}]
|
|
})
|
|
|
|
// Mock verb retrieval
|
|
s3Mock.on(GetObjectCommand, {
|
|
Key: 'verbs/verb-1.json'
|
|
}).resolves({
|
|
Body: createMockBody(verbData)
|
|
})
|
|
|
|
brainy = new BrainyData({
|
|
embeddingFunction: createMockEmbeddingFunction(),
|
|
storage: {
|
|
s3Storage: {
|
|
bucketName: 'test-bucket',
|
|
accessKeyId: 'test-key',
|
|
secretAccessKey: 'test-secret',
|
|
region: 'us-east-1'
|
|
}
|
|
}
|
|
})
|
|
await brainy.init()
|
|
|
|
// Create nouns first
|
|
const id1 = await brainy.add('Noun 1', { type: 'entity' })
|
|
const id2 = await brainy.add('Noun 2', { type: 'entity' })
|
|
|
|
// Create relationship
|
|
await brainy.relate(id1, id2, 'relates_to')
|
|
|
|
// Verify verb was saved
|
|
const putCalls = s3Mock.commandCalls(PutObjectCommand)
|
|
const verbSave = putCalls.find(call =>
|
|
call.args[0].input.Key?.includes('verbs/')
|
|
)
|
|
expect(verbSave).toBeDefined()
|
|
|
|
// Get verbs by source
|
|
const verbs = await brainy.getVerbsBySource(id1)
|
|
expect(verbs.length).toBeGreaterThanOrEqual(0) // May be 0 if not mocked fully
|
|
})
|
|
})
|
|
|
|
describe('Metadata', () => {
|
|
it('should handle metadata storage and retrieval', async () => {
|
|
const metadata = {
|
|
name: 'Test Item',
|
|
category: 'test',
|
|
tags: ['tag1', 'tag2'],
|
|
nested: {
|
|
property: 'value'
|
|
}
|
|
}
|
|
|
|
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 item with complex metadata
|
|
const id = await brainy.add('Test data', metadata)
|
|
|
|
// Update metadata
|
|
await brainy.updateMetadata(id, {
|
|
...metadata,
|
|
updated: true
|
|
})
|
|
|
|
// Verify metadata was saved
|
|
const putCalls = s3Mock.commandCalls(PutObjectCommand)
|
|
const metadataSave = putCalls.find(call =>
|
|
call.args[0].input.Key?.includes('metadata/')
|
|
)
|
|
expect(metadataSave).toBeDefined()
|
|
|
|
// Retrieve metadata
|
|
const retrieved = await brainy.getMetadata(id)
|
|
expect(retrieved).toBeDefined()
|
|
expect(retrieved?.name).toBe('Test Item')
|
|
})
|
|
|
|
it('should handle batch metadata operations', 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 multiple items
|
|
const ids = []
|
|
for (let i = 0; i < 10; i++) {
|
|
const id = await brainy.add(`Item ${i}`, { index: i })
|
|
ids.push(id)
|
|
}
|
|
|
|
// Batch get metadata
|
|
const metadataList = await brainy.getBatch(ids)
|
|
expect(metadataList.length).toBe(10)
|
|
})
|
|
})
|
|
|
|
describe('HNSW Index', () => {
|
|
it('should save and load HNSW index from S3', async () => {
|
|
const indexData = {
|
|
nodes: {
|
|
'node-1': {
|
|
id: 'node-1',
|
|
vector: [0.1, 0.2, 0.3],
|
|
connections: {
|
|
0: ['node-2', 'node-3']
|
|
},
|
|
level: 1
|
|
}
|
|
},
|
|
entryPoint: 'node-1',
|
|
dimensions: 3,
|
|
efConstruction: 200,
|
|
m: 16
|
|
}
|
|
|
|
// Mock index retrieval
|
|
s3Mock.on(GetObjectCommand, {
|
|
Key: 'index/hnsw-index.json'
|
|
}).resolves({
|
|
Body: createMockBody(indexData)
|
|
})
|
|
|
|
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 items to build index
|
|
await brainy.add('Vector 1 content', { name: 'Vector 1' })
|
|
await brainy.add('Vector 2 content', { name: 'Vector 2' })
|
|
|
|
// Search to verify index works
|
|
const results = await brainy.search('search query', 5)
|
|
expect(results).toBeDefined()
|
|
})
|
|
})
|
|
})
|
|
|
|
describe('S3 Error Handling', () => {
|
|
it('should handle S3 connection failures gracefully', async () => {
|
|
// Simulate S3 connection failure
|
|
s3Mock.on(GetObjectCommand).rejects(new Error('Connection timeout'))
|
|
s3Mock.on(PutObjectCommand).rejects(new Error('Connection timeout'))
|
|
|
|
brainy = new BrainyData({
|
|
storage: {
|
|
s3Storage: {
|
|
bucketName: 'test-bucket',
|
|
accessKeyId: 'test-key',
|
|
secretAccessKey: 'test-secret',
|
|
region: 'us-east-1'
|
|
}
|
|
}
|
|
})
|
|
|
|
// Should handle initialization failure gracefully
|
|
await expect(brainy.init()).resolves.not.toThrow()
|
|
})
|
|
|
|
it('should retry on transient S3 errors', async () => {
|
|
let attempts = 0
|
|
|
|
// Fail first 2 attempts, succeed on third
|
|
s3Mock.on(PutObjectCommand).callsFake(() => {
|
|
attempts++
|
|
if (attempts < 3) {
|
|
const error: any = new Error('Service Unavailable')
|
|
error.$metadata = { httpStatusCode: 503 }
|
|
throw error
|
|
}
|
|
return Promise.resolve({})
|
|
})
|
|
|
|
s3Mock.on(GetObjectCommand).rejects({ name: 'NoSuchKey' })
|
|
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()
|
|
|
|
// Should retry and eventually succeed
|
|
await brainy.add('Test data', { metadata: 'test' })
|
|
|
|
// Verify retries occurred
|
|
expect(attempts).toBeGreaterThanOrEqual(1)
|
|
})
|
|
|
|
it('should handle S3 permission errors', async () => {
|
|
// Simulate permission denied
|
|
const permissionError: any = new Error('Access Denied')
|
|
permissionError.name = 'AccessDenied'
|
|
permissionError.$metadata = { httpStatusCode: 403 }
|
|
|
|
s3Mock.on(GetObjectCommand).rejects(permissionError)
|
|
s3Mock.on(PutObjectCommand).rejects(permissionError)
|
|
|
|
brainy = new BrainyData({
|
|
storage: {
|
|
s3Storage: {
|
|
bucketName: 'test-bucket',
|
|
accessKeyId: 'invalid-key',
|
|
secretAccessKey: 'invalid-secret',
|
|
region: 'us-east-1'
|
|
}
|
|
}
|
|
})
|
|
|
|
// Should handle permission errors without crashing
|
|
await expect(brainy.init()).resolves.not.toThrow()
|
|
})
|
|
})
|
|
|
|
describe('S3 Performance Optimizations', () => {
|
|
it('should use multipart upload for large objects', 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()
|
|
|
|
// Create large dataset
|
|
const largeData = Array(1000).fill(0).map((_, i) => ({
|
|
data: `Large item ${i}`,
|
|
metadata: {
|
|
index: i,
|
|
largeField: 'x'.repeat(1000) // Make metadata large
|
|
}
|
|
}))
|
|
|
|
// Add large batch
|
|
await brainy.addBatch(
|
|
largeData.map(d => d.data),
|
|
largeData.map(d => d.metadata)
|
|
)
|
|
|
|
// Verify data was uploaded
|
|
const putCalls = s3Mock.commandCalls(PutObjectCommand)
|
|
expect(putCalls.length).toBeGreaterThan(0)
|
|
})
|
|
|
|
it('should implement caching to reduce S3 calls', async () => {
|
|
const testData = {
|
|
id: 'cached-item',
|
|
vector: [0.1, 0.2, 0.3],
|
|
metadata: { cached: true }
|
|
}
|
|
|
|
let getCalls = 0
|
|
s3Mock.on(GetObjectCommand).callsFake((input) => {
|
|
getCalls++
|
|
if (input.Key?.includes('nouns/cached-item')) {
|
|
return Promise.resolve({
|
|
Body: createMockBody(testData)
|
|
})
|
|
}
|
|
return Promise.reject({ 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'
|
|
},
|
|
cacheConfig: {
|
|
hotCacheMaxSize: 100,
|
|
warmCacheTTL: 60000
|
|
}
|
|
}
|
|
})
|
|
await brainy.init()
|
|
|
|
// Add item
|
|
const id = await brainy.add('Cached content', { cached: true }, { id: 'cached-item' })
|
|
|
|
// First get - should hit S3
|
|
await brainy.get(id)
|
|
const firstGetCalls = getCalls
|
|
|
|
// Second get - should hit cache
|
|
await brainy.get(id)
|
|
const secondGetCalls = getCalls
|
|
|
|
// Cache should prevent additional S3 call
|
|
expect(secondGetCalls).toBe(firstGetCalls)
|
|
})
|
|
|
|
it('should parallelize S3 operations for better throughput', 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 multiple items concurrently
|
|
const startTime = Date.now()
|
|
const promises = []
|
|
|
|
for (let i = 0; i < 20; i++) {
|
|
promises.push(
|
|
brainy.add(`Parallel item ${i}`, { index: i })
|
|
)
|
|
}
|
|
|
|
const ids = await Promise.all(promises)
|
|
const endTime = Date.now()
|
|
|
|
expect(ids.length).toBe(20)
|
|
|
|
// Operations should be parallelized (fast)
|
|
// In real scenario, this would be much faster than sequential
|
|
expect(endTime - startTime).toBeLessThan(5000) // Should be fast due to mocking
|
|
})
|
|
})
|
|
|
|
describe('S3 Data Integrity', () => {
|
|
it('should verify data integrity with checksums', async () => {
|
|
s3Mock.on(GetObjectCommand).rejects({ name: 'NoSuchKey' })
|
|
|
|
// Track checksums in PUT operations
|
|
const putChecksums: string[] = []
|
|
s3Mock.on(PutObjectCommand).callsFake((input) => {
|
|
if (input.ChecksumAlgorithm || input.ChecksumCRC32) {
|
|
putChecksums.push(input.Key || '')
|
|
}
|
|
return Promise.resolve({})
|
|
})
|
|
|
|
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
|
|
await brainy.add('Important data', { critical: true })
|
|
|
|
// Verify checksum was used for critical data
|
|
const putCalls = s3Mock.commandCalls(PutObjectCommand)
|
|
expect(putCalls.length).toBeGreaterThan(0)
|
|
})
|
|
|
|
it('should handle corrupted data gracefully', async () => {
|
|
// Return corrupted JSON
|
|
s3Mock.on(GetObjectCommand).resolves({
|
|
Body: {
|
|
transformToString: async () => '{ corrupt json ][',
|
|
transformToByteArray: async () => new TextEncoder().encode('{ corrupt json ][')
|
|
}
|
|
})
|
|
|
|
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'
|
|
}
|
|
}
|
|
})
|
|
|
|
// Should handle corrupted data without crashing
|
|
await expect(brainy.init()).resolves.not.toThrow()
|
|
})
|
|
})
|
|
|
|
describe('S3 Cleanup Operations', () => {
|
|
it('should properly clean up S3 objects on delete', async () => {
|
|
s3Mock.on(GetObjectCommand).rejects({ name: 'NoSuchKey' })
|
|
s3Mock.on(PutObjectCommand).resolves({})
|
|
s3Mock.on(DeleteObjectCommand).resolves({})
|
|
s3Mock.on(DeleteObjectsCommand).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 and then delete item
|
|
const id = await brainy.add('To be deleted', { temporary: true })
|
|
await brainy.delete(id)
|
|
|
|
// Verify delete was called
|
|
const deleteCalls = s3Mock.commandCalls(DeleteObjectCommand)
|
|
expect(deleteCalls.length).toBeGreaterThan(0)
|
|
})
|
|
|
|
it('should batch delete operations for efficiency', async () => {
|
|
s3Mock.on(GetObjectCommand).rejects({ name: 'NoSuchKey' })
|
|
s3Mock.on(PutObjectCommand).resolves({})
|
|
s3Mock.on(DeleteObjectsCommand).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()
|
|
|
|
// Clear all data
|
|
await brainy.clear()
|
|
|
|
// Verify batch delete was used
|
|
const batchDeleteCalls = s3Mock.commandCalls(DeleteObjectsCommand)
|
|
// Clear operation should use batch delete
|
|
expect(batchDeleteCalls.length).toBeGreaterThanOrEqual(0)
|
|
})
|
|
})
|
|
}) |