brainy/tests/augmentations-entity-registry.test.ts

import { describe, it, expect, beforeEach, afterEach } from 'vitest'
import { BrainyData } from '../src/index.js'
import { EntityRegistryAugmentation, AutoRegisterEntitiesAugmentation } from '../src/augmentations/entityRegistryAugmentation.js'

describe('Entity Registry Augmentation', () => {
  let db: BrainyData | null = null
  
  // Helper to create test vectors
  const createTestVector = (seed: number = 0) => {
    return new Array(384).fill(0).map((_, i) => Math.sin(i + seed) * 0.5)
  }
  
  afterEach(async () => {
    if (db) {
      await db.cleanup?.()
      db = null
    }
    
    // Force garbage collection if available
    if (global.gc) {
      global.gc()
    }
  })
  
  describe('Entity Registry - Fast Deduplication', () => {
    beforeEach(async () => {
      db = new BrainyData({
        entityRegistry: {
          enabled: true,
          maxCacheSize: 1000,
          cacheTTL: 5000, // 5 seconds
          persistence: 'memory',
          indexedFields: ['did', 'handle', 'uri', 'external_id', 'id']
        }
      })
      await db.init()
    })
    
    it('should register entities automatically', async () => {
      // Add entity with external ID
      await db.add(createTestVector(1), { 
        id: 'internal1',
        external_id: 'ext123',
        data: 'Test entity'
      })
      
      // Should be able to lookup by external ID quickly
      const result = await db.get('internal1')
      expect(result).toBeDefined()
      expect(result?.metadata?.external_id).toBe('ext123')
    })
    
    it('should prevent duplicate entities', async () => {
      // Add first entity
      const id1 = await db.add(createTestVector(1), { 
        external_id: 'unique123',
        data: 'Original'
      })
      
      // Try to add duplicate with same external_id
      const id2 = await db.add(createTestVector(2), { 
        external_id: 'unique123',
        data: 'Duplicate attempt'
      })
      
      // Should return same ID (deduplicated)
      expect(id1).toBe(id2)
      
      // Data should be from original
      const result = await db.get(id1)
      expect(result?.metadata?.data).toBe('Original')
    })
    
    it('should handle high-throughput streaming data', async () => {
      const startTime = performance.now()
      const promises = []
      
      // Simulate streaming data with some duplicates
      for (let i = 0; i < 1000; i++) {
        const externalId = `stream${i % 500}` // 50% duplicates
        promises.push(
          db!.add(createTestVector(i), {
            external_id: externalId,
            data: `Stream item ${i}`
          })
        )
      }
      
      const ids = await Promise.all(promises)
      const uniqueIds = new Set(ids)
      
      // Should have deduplicated to ~500 unique items
      expect(uniqueIds.size).toBeLessThanOrEqual(500)
      
      const elapsed = performance.now() - startTime
      // Should be fast (< 2 seconds for 1000 items)
      expect(elapsed).toBeLessThan(2000)
    })
    
    it('should support multiple indexed fields', async () => {
      // Add entity with multiple identifiers
      await db.add(createTestVector(1), {
        id: 'internal1',
        did: 'did:example:123',
        handle: '@user.example',
        uri: 'https://example.com/user',
        external_id: 'ext456',
        data: 'Multi-ID entity'
      })
      
      // Should be findable by any indexed field
      // (Note: actual lookup by these fields would need specific API methods)
      const result = await db.get('internal1')
      expect(result?.metadata?.did).toBe('did:example:123')
      expect(result?.metadata?.handle).toBe('@user.example')
      expect(result?.metadata?.uri).toBe('https://example.com/user')
    })
    
    it('should handle cache expiration', async () => {
      const registry = new EntityRegistryAugmentation({
        maxCacheSize: 10,
        cacheTTL: 100, // 100ms TTL for testing
        persistence: 'memory'
      })
      
      // Initialize with mock context
      await registry.initialize({
        brain: db,
        storage: {},
        config: {},
        log: () => {}
      } as any)
      
      // Register an entity
      const registered = registry.register('ext789', 'internal789')
      expect(registered).toBe(true)
      
      // Should be in cache
      const immediate = registry.lookup('ext789')
      expect(immediate).toBe('internal789')
      
      // Wait for TTL to expire
      await new Promise(resolve => setTimeout(resolve, 150))
      
      // Should be expired from cache
      const expired = registry.lookup('ext789')
      expect(expired).toBeNull()
    })
  })
  
  describe('Auto-Register Entities', () => {
    beforeEach(async () => {
      db = new BrainyData({
        entityRegistry: {
          enabled: true
        },
        autoRegisterEntities: {
          enabled: true
        }
      })
      await db.init()
    })
    
    it('should auto-register entities after adding', async () => {
      // Add entity
      const id = await db.add(createTestVector(1), {
        external_id: 'auto123',
        handle: '@auto.user',
        data: 'Auto-registered'
      })
      
      // Should be registered automatically
      const result = await db.get(id)
      expect(result).toBeDefined()
      expect(result?.metadata?.external_id).toBe('auto123')
    })
    
    it('should work with batch operations', async () => {
      const items = []
      for (let i = 0; i < 100; i++) {
        items.push({
          vector: createTestVector(i),
          metadata: {
            external_id: `batch${i}`,
            data: `Batch item ${i}`
          }
        })
      }
      
      // Add batch
      const ids = await Promise.all(
        items.map(item => db!.add(item.vector, item.metadata))
      )
      
      // All should be registered
      expect(ids.length).toBe(100)
      const uniqueIds = new Set(ids)
      expect(uniqueIds.size).toBe(100) // All unique
    })
  })
  
  describe('Performance Benchmarks', () => {
    it('should provide O(1) lookup performance', async () => {
      db = new BrainyData({
        entityRegistry: {
          enabled: true,
          maxCacheSize: 100000
        }
      })
      await db.init()
      
      // Add many entities
      for (let i = 0; i < 10000; i++) {
        await db.add(createTestVector(i), {
          id: `perf${i}`,
          external_id: `ext${i}`
        })
      }
      
      // Measure lookup time
      const lookupTimes = []
      for (let i = 0; i < 100; i++) {
        const randomId = Math.floor(Math.random() * 10000)
        const start = performance.now()
        await db.get(`perf${randomId}`)
        lookupTimes.push(performance.now() - start)
      }
      
      // Average lookup should be very fast (< 1ms)
      const avgLookup = lookupTimes.reduce((a, b) => a + b, 0) / lookupTimes.length
      expect(avgLookup).toBeLessThan(1)
    })
    
    it('should handle cache size limits efficiently', async () => {
      const registry = new EntityRegistryAugmentation({
        maxCacheSize: 100, // Small cache
        cacheTTL: 60000,
        persistence: 'memory'
      })
      
      await registry.initialize({
        brain: {},
        storage: {},
        config: {},
        log: () => {}
      } as any)
      
      // Add more than cache size
      for (let i = 0; i < 200; i++) {
        registry.register(`ext${i}`, `internal${i}`)
      }
      
      // Cache should not exceed max size
      const stats = registry.getStats()
      expect(stats.cacheSize).toBeLessThanOrEqual(100)
      expect(stats.totalRegistered).toBe(200)
    })
  })
  
  describe('Persistence Options', () => {
    it('should support memory persistence', async () => {
      const registry = new EntityRegistryAugmentation({
        persistence: 'memory'
      })
      
      await registry.initialize({
        brain: {},
        storage: {},
        config: {},
        log: () => {}
      } as any)
      
      registry.register('mem1', 'internal1')
      expect(registry.lookup('mem1')).toBe('internal1')
      
      // Memory persistence doesn't survive restart
      const newRegistry = new EntityRegistryAugmentation({
        persistence: 'memory'
      })
      
      await newRegistry.initialize({
        brain: {},
        storage: {},
        config: {},
        log: () => {}
      } as any)
      
      expect(newRegistry.lookup('mem1')).toBeNull()
    })
    
    it('should support hybrid persistence', async () => {
      const registry = new EntityRegistryAugmentation({
        persistence: 'hybrid',
        maxCacheSize: 50
      })
      
      await registry.initialize({
        brain: {},
        storage: {
          saveMetadata: async () => {},
          getMetadata: async () => null
        },
        config: {},
        log: () => {}
      } as any)
      
      // Add many items
      for (let i = 0; i < 100; i++) {
        registry.register(`hybrid${i}`, `internal${i}`)
      }
      
      const stats = registry.getStats()
      // Hot cache should be limited
      expect(stats.cacheSize).toBeLessThanOrEqual(50)
      // But all should be registered
      expect(stats.totalRegistered).toBe(100)
    })
  })
  
  describe('Standalone Usage', () => {
    it('should work as standalone augmentation', () => {
      const registry = new EntityRegistryAugmentation({
        maxCacheSize: 1000
      })
      
      expect(registry.name).toBe('EntityRegistry')
      expect(registry.timing).toBe('before')
      expect(registry.priority).toBe(95) // High priority
      expect(registry.operations).toContain('add')
      expect(registry.operations).toContain('addNoun')
    })
    
    it('should provide comprehensive statistics', async () => {
      const registry = new EntityRegistryAugmentation()
      
      await registry.initialize({
        brain: {},
        storage: {},
        config: {},
        log: () => {}
      } as any)
      
      // Register some entities
      registry.register('ext1', 'int1')
      registry.register('ext2', 'int2')
      registry.register('ext1', 'int1') // Duplicate
      
      const stats = registry.getStats()
      expect(stats.totalRegistered).toBe(2)
      expect(stats.cacheSize).toBe(2)
      expect(stats.hits).toBe(0)
      expect(stats.misses).toBe(0)
      
      // Lookup to generate hits/misses
      registry.lookup('ext1') // Hit
      registry.lookup('ext3') // Miss
      
      const newStats = registry.getStats()
      expect(newStats.hits).toBe(1)
      expect(newStats.misses).toBe(1)
    })
  })
  
  describe('Error Handling', () => {
    it('should handle invalid external IDs', async () => {
      db = new BrainyData({
        entityRegistry: { enabled: true }
      })
      await db.init()
      
      // Should handle null/undefined external IDs
      const id1 = await db.add(createTestVector(1), {
        external_id: null,
        data: 'No external ID'
      })
      
      expect(id1).toBeDefined()
      
      // Should handle empty string
      const id2 = await db.add(createTestVector(2), {
        external_id: '',
        data: 'Empty external ID'
      })
      
      expect(id2).toBeDefined()
      expect(id2).not.toBe(id1) // Should be different
    })
    
    it('should handle registration failures gracefully', () => {
      const registry = new EntityRegistryAugmentation()
      
      // Try to use before initialization
      const result = registry.register('test', 'test')
      expect(result).toBe(false) // Should fail gracefully
      
      const lookup = registry.lookup('test')
      expect(lookup).toBeNull()
    })
  })
})