@dotdo/postgres 0.1.0

package/src/pglite/tiered-init.test.ts

@@ -0,0 +1,725 @@
+/**
+ * Tests for Tiered Storage PGLite Initialization
+ *
+ * Tests cover:
+ * - Full tiered initialization flow
+ * - Tier selection logic
+ * - Metrics collection
+ * - Callbacks and hooks
+ * - Error handling
+ * - Utility functions
+ */
+
+import { describe, it, expect, vi, beforeEach } from 'vitest'
+import {
+  initWithTieredStorage,
+  initWithTier,
+  createTieredInitializer,
+  determineBestTier,
+  checkTierAvailability,
+  persistToTier,
+  type TieredInitConfig,
+} from './tiered-init'
+import type { CacheLayer } from './cache-layer'
+import type { R2StorageLayer } from './r2-layer'
+
+// Mock PGLite
+vi.mock('@dotdo/pglite', () => {
+  return {
+    PGlite: {
+      create: vi.fn().mockImplementation(async (options: unknown) => {
+        return {
+          query: vi.fn().mockResolvedValue({ rows: [], fields: [] }),
+          exec: vi.fn().mockResolvedValue([]),
+          close: vi.fn().mockResolvedValue(undefined),
+          isReady: true,
+          waitReady: Promise.resolve(),
+          dumpDataDir: vi.fn().mockResolvedValue(new Uint8Array([1, 2, 3])),
+        }
+      }),
+    },
+  }
+})
+
+// ============================================================================
+// Mock Factories
+// ============================================================================
+
+function createMockWasmModule(): WebAssembly.Module {
+  return {} as WebAssembly.Module
+}
+
+function createMockCacheLayer(data?: Uint8Array): CacheLayer {
+  return {
+    get: vi.fn().mockResolvedValue(data || null),
+    put: vi.fn().mockResolvedValue(undefined),
+    delete: vi.fn().mockResolvedValue(true),
+    has: vi.fn().mockResolvedValue(!!data),
+    getWithMetadata: vi.fn().mockResolvedValue(data ? { data, metadata: {} } : null),
+    clear: vi.fn().mockResolvedValue(undefined),
+    getStats: vi.fn().mockReturnValue({
+      hits: 0,
+      misses: 0,
+      writes: 0,
+      deletes: 0,
+      bytesRead: 0,
+      bytesWritten: 0,
+      errors: 0,
+      hitRatio: 0,
+    }),
+  } as unknown as CacheLayer
+}
+
+function createMockR2Layer(data?: Uint8Array): R2StorageLayer {
+  return {
+    get: vi.fn().mockResolvedValue(data || null),
+    put: vi.fn().mockResolvedValue(undefined),
+    delete: vi.fn().mockResolvedValue(undefined),
+    has: vi.fn().mockResolvedValue(!!data),
+    getRange: vi.fn().mockResolvedValue(data || null),
+    head: vi.fn().mockResolvedValue(
+      data ? { size: data.length, etag: 'test', uploaded: new Date() } : null
+    ),
+    list: vi.fn().mockResolvedValue({ keys: [], truncated: false }),
+    deleteMany: vi.fn().mockResolvedValue(undefined),
+    getStats: vi.fn().mockReturnValue({
+      reads: 0,
+      rangeReads: 0,
+      writes: 0,
+      deletes: 0,
+      bytesRead: 0,
+      bytesWritten: 0,
+      errors: 0,
+    }),
+  } as unknown as R2StorageLayer
+}
+
+function createMockR2Bucket(data?: Uint8Array): R2Bucket {
+  return {
+    get: vi.fn().mockResolvedValue(
+      data
+        ? {
+            arrayBuffer: vi.fn().mockResolvedValue(data.buffer),
+            bytes: vi.fn().mockResolvedValue(data),
+          }
+        : null
+    ),
+    put: vi.fn().mockResolvedValue(undefined),
+    delete: vi.fn().mockResolvedValue(undefined),
+    head: vi.fn().mockResolvedValue(data ? { size: data.length } : null),
+    list: vi.fn().mockResolvedValue({ objects: [], truncated: false }),
+  } as unknown as R2Bucket
+}
+
+function createMockCache(data?: ArrayBuffer): Cache {
+  return {
+    match: vi.fn().mockResolvedValue(
+      data
+        ? {
+            arrayBuffer: vi.fn().mockResolvedValue(data),
+          }
+        : undefined
+    ),
+    put: vi.fn().mockResolvedValue(undefined),
+    delete: vi.fn().mockResolvedValue(true),
+  } as unknown as Cache
+}
+
+// ============================================================================
+// Tests
+// ============================================================================
+
+describe('initWithTieredStorage', () => {
+  describe('Hot Data Path', () => {
+    it('should initialize with hot data', async () => {
+      const hotData = new Uint8Array([1, 2, 3, 4, 5])
+      const wasmModule = createMockWasmModule()
+
+      const { pglite, metrics } = await initWithTieredStorage({
+        wasmModule,
+        hotData,
+      })
+
+      expect(pglite).toBeDefined()
+      expect(metrics.tier).toBe('hot')
+      expect(metrics.dataSize).toBe(5)
+    })
+
+    it('should skip warm/cold when hot data is provided', async () => {
+      const hotData = new Uint8Array([1, 2, 3])
+      const wasmModule = createMockWasmModule()
+      const cacheLayer = createMockCacheLayer(new Uint8Array([4, 5, 6]))
+      const r2Layer = createMockR2Layer(new Uint8Array([7, 8, 9]))
+
+      const { metrics } = await initWithTieredStorage({
+        wasmModule,
+        hotData,
+        cacheLayer,
+        warmCacheKey: 'test',
+        r2Layer,
+        coldStorageKey: 'test',
+      })
+
+      expect(metrics.tier).toBe('hot')
+      expect(cacheLayer.get).not.toHaveBeenCalled()
+      expect(r2Layer.get).not.toHaveBeenCalled()
+    })
+  })
+
+  describe('Warm Cache Path', () => {
+    it('should initialize from cache layer', async () => {
+      const cachedData = new Uint8Array([1, 2, 3])
+      const wasmModule = createMockWasmModule()
+      const cacheLayer = createMockCacheLayer(cachedData)
+
+      const { pglite, metrics } = await initWithTieredStorage({
+        wasmModule,
+        cacheLayer,
+        warmCacheKey: 'test-key',
+      })
+
+      expect(pglite).toBeDefined()
+      expect(metrics.tier).toBe('warm')
+      expect(metrics.dataSize).toBe(3)
+    })
+  })
+
+  describe('Cold Storage Path', () => {
+    it('should initialize from R2 layer', async () => {
+      const coldData = new Uint8Array([1, 2, 3, 4])
+      const wasmModule = createMockWasmModule()
+      const r2Layer = createMockR2Layer(coldData)
+
+      const { pglite, metrics } = await initWithTieredStorage({
+        wasmModule,
+        r2Layer,
+        coldStorageKey: 'test-key',
+      })
+
+      expect(pglite).toBeDefined()
+      expect(metrics.tier).toBe('cold')
+      expect(metrics.dataSize).toBe(4)
+    })
+
+    it('should initialize from R2 bucket', async () => {
+      const coldData = new Uint8Array([1, 2, 3])
+      const wasmModule = createMockWasmModule()
+      const coldBucket = createMockR2Bucket(coldData)
+
+      const { pglite, metrics } = await initWithTieredStorage({
+        wasmModule,
+        coldBucket,
+        coldStorageKey: 'test-key',
+      })
+
+      expect(pglite).toBeDefined()
+      expect(metrics.tier).toBe('cold')
+    })
+  })
+
+  describe('Metrics', () => {
+    it('should track timing metrics', async () => {
+      const wasmModule = createMockWasmModule()
+
+      const { metrics } = await initWithTieredStorage({
+        wasmModule,
+      })
+
+      expect(metrics.totalMs).toBeGreaterThanOrEqual(0)
+      expect(metrics.dataLoadMs).toBeGreaterThanOrEqual(0)
+      expect(metrics.pgliteCreateMs).toBeGreaterThanOrEqual(0)
+      expect(metrics.startedAt).toBeGreaterThan(0)
+      expect(metrics.completedAt).toBeGreaterThanOrEqual(metrics.startedAt)
+    })
+
+    it('should track parallelization status', async () => {
+      const wasmModule = createMockWasmModule()
+      const cacheLayer = createMockCacheLayer()
+      const r2Layer = createMockR2Layer(new Uint8Array([1, 2, 3]))
+
+      const { metrics } = await initWithTieredStorage({
+        wasmModule,
+        cacheLayer,
+        warmCacheKey: 'warm',
+        r2Layer,
+        coldStorageKey: 'cold',
+      })
+
+      expect(metrics.parallelized).toBe(true)
+    })
+  })
+
+  describe('Callbacks', () => {
+    it('should call onTierHit callback', async () => {
+      const coldData = new Uint8Array([1, 2, 3])
+      const wasmModule = createMockWasmModule()
+      const r2Layer = createMockR2Layer(coldData)
+      const onTierHit = vi.fn()
+
+      await initWithTieredStorage({
+        wasmModule,
+        r2Layer,
+        coldStorageKey: 'test-key',
+        onTierHit,
+      })
+
+      expect(onTierHit).toHaveBeenCalledWith('cold', 3, expect.any(Number))
+    })
+
+    it('should not call onTierHit when no data found', async () => {
+      const wasmModule = createMockWasmModule()
+      const onTierHit = vi.fn()
+
+      await initWithTieredStorage({
+        wasmModule,
+        onTierHit,
+      })
+
+      expect(onTierHit).not.toHaveBeenCalled()
+    })
+  })
+
+  describe('Empty Data', () => {
+    it('should handle empty data gracefully', async () => {
+      const wasmModule = createMockWasmModule()
+
+      const { pglite, metrics } = await initWithTieredStorage({
+        wasmModule,
+      })
+
+      expect(pglite).toBeDefined()
+      expect(metrics.tier).toBe('empty')
+      expect(metrics.dataSize).toBe(0)
+    })
+  })
+
+  describe('PGLite Options', () => {
+    it('should pass custom database name', async () => {
+      const wasmModule = createMockWasmModule()
+      const { PGlite } = await import('@dotdo/pglite')
+
+      await initWithTieredStorage({
+        wasmModule,
+        database: 'custom-db',
+      })
+
+      expect(PGlite.create).toHaveBeenCalledWith(
+        expect.objectContaining({
+          database: 'custom-db',
+        })
+      )
+    })
+
+    it('should pass additional pglite options', async () => {
+      const wasmModule = createMockWasmModule()
+      const { PGlite } = await import('@dotdo/pglite')
+
+      await initWithTieredStorage({
+        wasmModule,
+        pgliteOptions: {
+          debug: true,
+        },
+      })
+
+      expect(PGlite.create).toHaveBeenCalledWith(
+        expect.objectContaining({
+          debug: true,
+        })
+      )
+    })
+  })
+})
+
+describe('initWithTier', () => {
+  describe('Hot Tier', () => {
+    it('should initialize with hot data', async () => {
+      const hotData = new ArrayBuffer(10)
+      const wasmModule = createMockWasmModule()
+
+      const pglite = await initWithTier({
+        wasmModule,
+        tier: 'hot',
+        hotData,
+        dataKey: 'test',
+      })
+
+      expect(pglite).toBeDefined()
+    })
+
+    it('should throw error when hot data is missing', async () => {
+      const wasmModule = createMockWasmModule()
+
+      await expect(
+        initWithTier({
+          wasmModule,
+          tier: 'hot',
+          dataKey: 'test',
+        })
+      ).rejects.toThrow('Hot tier requires hotData')
+    })
+  })
+
+  describe('Warm Tier', () => {
+    it('should initialize from warm cache', async () => {
+      const cachedData = new ArrayBuffer(10)
+      const wasmModule = createMockWasmModule()
+      const warmCache = createMockCache(cachedData)
+
+      const pglite = await initWithTier({
+        wasmModule,
+        tier: 'warm',
+        warmCache,
+        dataKey: 'test-key',
+      })
+
+      expect(pglite).toBeDefined()
+      expect(warmCache.match).toHaveBeenCalledWith('test-key')
+    })
+
+    it('should throw error when warm cache is missing', async () => {
+      const wasmModule = createMockWasmModule()
+
+      await expect(
+        initWithTier({
+          wasmModule,
+          tier: 'warm',
+          dataKey: 'test',
+        })
+      ).rejects.toThrow('Warm tier requires warmCache')
+    })
+
+    it('should handle cache miss', async () => {
+      const wasmModule = createMockWasmModule()
+      const warmCache = createMockCache() // No data
+
+      const pglite = await initWithTier({
+        wasmModule,
+        tier: 'warm',
+        warmCache,
+        dataKey: 'test-key',
+      })
+
+      expect(pglite).toBeDefined()
+    })
+  })
+
+  describe('Cold Tier', () => {
+    it('should initialize from cold bucket', async () => {
+      const coldData = new Uint8Array([1, 2, 3])
+      const wasmModule = createMockWasmModule()
+      const coldBucket = createMockR2Bucket(coldData)
+
+      const pglite = await initWithTier({
+        wasmModule,
+        tier: 'cold',
+        coldBucket,
+        dataKey: 'test-key',
+      })
+
+      expect(pglite).toBeDefined()
+      expect(coldBucket.get).toHaveBeenCalledWith('test-key')
+    })
+
+    it('should throw error when cold bucket is missing', async () => {
+      const wasmModule = createMockWasmModule()
+
+      await expect(
+        initWithTier({
+          wasmModule,
+          tier: 'cold',
+          dataKey: 'test',
+        })
+      ).rejects.toThrow('Cold tier requires coldBucket')
+    })
+  })
+
+  describe('Invalid Tier', () => {
+    it('should throw error for unknown tier', async () => {
+      const wasmModule = createMockWasmModule()
+
+      await expect(
+        initWithTier({
+          wasmModule,
+          tier: 'invalid' as 'hot',
+          dataKey: 'test',
+        })
+      ).rejects.toThrow('Unknown tier')
+    })
+  })
+})
+
+describe('createTieredInitializer', () => {
+  it('should create a reusable initializer', async () => {
+    const wasmModule = createMockWasmModule()
+    const cacheLayer = createMockCacheLayer(new Uint8Array([1, 2, 3]))
+
+    const initPGLite = createTieredInitializer({
+      wasmModule,
+      cacheLayer,
+    })
+
+    const { pglite, metrics } = await initPGLite({
+      warmCacheKey: 'test-key',
+    })
+
+    expect(pglite).toBeDefined()
+    expect(metrics.tier).toBe('warm')
+  })
+
+  it('should allow different instance configs', async () => {
+    const wasmModule = createMockWasmModule()
+    const cacheLayer = createMockCacheLayer(new Uint8Array([1, 2, 3]))
+
+    const initPGLite = createTieredInitializer({
+      wasmModule,
+      cacheLayer,
+    })
+
+    const result1 = await initPGLite({ warmCacheKey: 'db1' })
+    const result2 = await initPGLite({ warmCacheKey: 'db2' })
+
+    expect(cacheLayer.get).toHaveBeenCalledWith('db1')
+    expect(cacheLayer.get).toHaveBeenCalledWith('db2')
+  })
+
+  it('should support onTierHit callback in instance config', async () => {
+    const wasmModule = createMockWasmModule()
+    const r2Layer = createMockR2Layer(new Uint8Array([1, 2, 3]))
+
+    const initPGLite = createTieredInitializer({
+      wasmModule,
+      r2Layer,
+    })
+
+    const onTierHit = vi.fn()
+    await initPGLite({
+      coldStorageKey: 'test-key',
+      onTierHit,
+    })
+
+    expect(onTierHit).toHaveBeenCalled()
+  })
+})
+
+describe('determineBestTier', () => {
+  it('should return hot when hot data is available', () => {
+    const tier = determineBestTier({
+      hasHotData: true,
+      hasCacheLayer: true,
+      hasR2Layer: true,
+    })
+
+    expect(tier).toBe('hot')
+  })
+
+  it('should return warm when cache layer is available', () => {
+    const tier = determineBestTier({
+      hasCacheLayer: true,
+      hasR2Layer: true,
+    })
+
+    expect(tier).toBe('warm')
+  })
+
+  it('should return warm when warm cache is available', () => {
+    const tier = determineBestTier({
+      hasWarmCache: true,
+      hasR2Layer: true,
+    })
+
+    expect(tier).toBe('warm')
+  })
+
+  it('should return cold when R2 layer is available', () => {
+    const tier = determineBestTier({
+      hasR2Layer: true,
+    })
+
+    expect(tier).toBe('cold')
+  })
+
+  it('should return cold when cold bucket is available', () => {
+    const tier = determineBestTier({
+      hasColdBucket: true,
+    })
+
+    expect(tier).toBe('cold')
+  })
+
+  it('should return empty when nothing is available', () => {
+    const tier = determineBestTier({})
+
+    expect(tier).toBe('empty')
+  })
+})
+
+describe('checkTierAvailability', () => {
+  it('should check warm tier availability', async () => {
+    const cacheLayer = createMockCacheLayer(new Uint8Array([1, 2, 3]))
+
+    const result = await checkTierAvailability({
+      cacheLayer,
+      warmCacheKey: 'test-key',
+    })
+
+    expect(result.warm).toBe(true)
+    expect(result.cold).toBe(false)
+    expect(result.recommendedTier).toBe('warm')
+  })
+
+  it('should check cold tier availability with R2 layer', async () => {
+    const r2Layer = createMockR2Layer(new Uint8Array([1, 2, 3]))
+
+    const result = await checkTierAvailability({
+      r2Layer,
+      coldStorageKey: 'test-key',
+    })
+
+    expect(result.warm).toBe(false)
+    expect(result.cold).toBe(true)
+    expect(result.recommendedTier).toBe('cold')
+  })
+
+  it('should check cold tier availability with R2 bucket', async () => {
+    const coldBucket = createMockR2Bucket(new Uint8Array([1, 2, 3]))
+
+    const result = await checkTierAvailability({
+      coldBucket,
+      coldStorageKey: 'test-key',
+    })
+
+    expect(result.cold).toBe(true)
+    expect(coldBucket.head).toHaveBeenCalledWith('test-key')
+  })
+
+  it('should check both tiers in parallel', async () => {
+    const cacheLayer = createMockCacheLayer(new Uint8Array([1, 2, 3]))
+    const r2Layer = createMockR2Layer(new Uint8Array([4, 5, 6]))
+
+    const result = await checkTierAvailability({
+      cacheLayer,
+      warmCacheKey: 'warm-key',
+      r2Layer,
+      coldStorageKey: 'cold-key',
+    })
+
+    expect(result.warm).toBe(true)
+    expect(result.cold).toBe(true)
+    expect(result.recommendedTier).toBe('warm') // Warm is preferred
+  })
+
+  it('should return empty when no data found', async () => {
+    const cacheLayer = createMockCacheLayer() // No data
+    const r2Layer = createMockR2Layer() // No data
+
+    const result = await checkTierAvailability({
+      cacheLayer,
+      warmCacheKey: 'warm-key',
+      r2Layer,
+      coldStorageKey: 'cold-key',
+    })
+
+    expect(result.warm).toBe(false)
+    expect(result.cold).toBe(false)
+    expect(result.recommendedTier).toBe('empty')
+  })
+})
+
+describe('persistToTier', () => {
+  describe('Warm Tier', () => {
+    it('should persist to cache layer', async () => {
+      const data = new Uint8Array([1, 2, 3])
+      const cacheLayer = createMockCacheLayer()
+
+      await persistToTier({
+        tier: 'warm',
+        data,
+        cacheLayer,
+        warmCacheKey: 'test-key',
+      })
+
+      expect(cacheLayer.put).toHaveBeenCalledWith('test-key', data, expect.any(Object))
+    })
+
+    it('should handle ArrayBuffer data', async () => {
+      const data = new ArrayBuffer(10)
+      const cacheLayer = createMockCacheLayer()
+
+      await persistToTier({
+        tier: 'warm',
+        data,
+        cacheLayer,
+        warmCacheKey: 'test-key',
+      })
+
+      expect(cacheLayer.put).toHaveBeenCalled()
+    })
+
+    it('should throw when cache layer is missing', async () => {
+      await expect(
+        persistToTier({
+          tier: 'warm',
+          data: new Uint8Array([1, 2, 3]),
+          warmCacheKey: 'test-key',
+        })
+      ).rejects.toThrow('Warm tier persistence requires cacheLayer and warmCacheKey')
+    })
+  })
+
+  describe('Cold Tier', () => {
+    it('should persist to R2 layer', async () => {
+      const data = new Uint8Array([1, 2, 3])
+      const r2Layer = createMockR2Layer()
+
+      await persistToTier({
+        tier: 'cold',
+        data,
+        r2Layer,
+        coldStorageKey: 'test-key',
+      })
+
+      expect(r2Layer.put).toHaveBeenCalledWith('test-key', data)
+    })
+
+    it('should persist to R2 bucket', async () => {
+      const data = new Uint8Array([1, 2, 3])
+      const coldBucket = createMockR2Bucket()
+
+      await persistToTier({
+        tier: 'cold',
+        data,
+        coldBucket,
+        coldStorageKey: 'test-key',
+      })
+
+      expect(coldBucket.put).toHaveBeenCalledWith('test-key', data)
+    })
+
+    it('should prefer R2 layer over bucket', async () => {
+      const data = new Uint8Array([1, 2, 3])
+      const r2Layer = createMockR2Layer()
+      const coldBucket = createMockR2Bucket()
+
+      await persistToTier({
+        tier: 'cold',
+        data,
+        r2Layer,
+        coldBucket,
+        coldStorageKey: 'test-key',
+      })
+
+      expect(r2Layer.put).toHaveBeenCalled()
+      expect(coldBucket.put).not.toHaveBeenCalled()
+    })
+
+    it('should throw when R2 is missing', async () => {
+      await expect(
+        persistToTier({
+          tier: 'cold',
+          data: new Uint8Array([1, 2, 3]),
+          coldStorageKey: 'test-key',
+        })
+      ).rejects.toThrow('Cold tier persistence requires')
+    })
+  })
+})