shieldx/tests/unit/learning/PatternStore.test.ts

/**
 * PatternStore tests — exercises the in-memory backend path (no DB required).
 * Validates pattern CRUD, incident tracking, stats, and deduplication.
 */
import { describe, it, expect, beforeEach } from 'vitest'
import { PatternStore } from '../../../src/learning/PatternStore.js'
import type { PatternRecord } from '../../../src/types/learning.js'
import type { ShieldXResult } from '../../../src/types/detection.js'

function makePattern(overrides: Partial<PatternRecord> = {}): PatternRecord {
  return {
    id: `pat-${Date.now()}-${Math.random()}`,
    createdAt: new Date().toISOString(),
    updatedAt: new Date().toISOString(),
    patternText: 'ignore all previous instructions',
    patternType: 'rule',
    killChainPhase: 'initial_access',
    confidenceBase: 0.9,
    hitCount: 0,
    falsePositiveCount: 0,
    source: 'builtin',
    enabled: true,
    ...overrides,
  }
}

function makeScanResult(overrides: Partial<ShieldXResult> = {}): ShieldXResult {
  return {
    id: `scan-${Date.now()}-${Math.random()}`,
    timestamp: new Date().toISOString(),
    input: 'test input',
    detected: true,
    threatLevel: 'high',
    killChainPhase: 'initial_access',
    action: 'block',
    scanResults: [],
    healingApplied: false,
    latencyMs: 5,
    ...overrides,
  }
}

describe('PatternStore (in-memory backend)', () => {
  let store: PatternStore

  beforeEach(async () => {
    store = new PatternStore({ backend: 'memory' })
    await store.initialize()
  })

  describe('initialize()', () => {
    it('should initialize without throwing', async () => {
      const s = new PatternStore({ backend: 'memory' })
      await expect(s.initialize()).resolves.not.toThrow()
    })

    it('should be idempotent on multiple calls', async () => {
      await expect(store.initialize()).resolves.not.toThrow()
      await expect(store.initialize()).resolves.not.toThrow()
    })
  })

  describe('savePattern() / loadPatterns()', () => {
    it('should save and retrieve a pattern', async () => {
      const pattern = makePattern({ id: 'test-001', patternText: 'ignore all previous' })
      await store.savePattern(pattern)

      const patterns = await store.loadPatterns()
      expect(patterns.length).toBeGreaterThan(0)
      const found = patterns.find((p) => p.id === 'test-001')
      expect(found).toBeDefined()
      expect(found!.patternText).toBe('ignore all previous')
    })

    it('should save multiple patterns', async () => {
      for (let i = 0; i < 5; i++) {
        await store.savePattern(
          makePattern({
            id: `pattern-${i}`,
            patternText: `test pattern ${i}`,
            confidenceBase: 0.8 + i * 0.02,
            hitCount: i,
          }),
        )
      }
      const patterns = await store.loadPatterns()
      expect(patterns.length).toBeGreaterThanOrEqual(5)
    })

    it('should update an existing pattern when saved with same id', async () => {
      await store.savePattern(
        makePattern({ id: 'update-test', patternText: 'original', confidenceBase: 0.5 }),
      )
      await store.savePattern(
        makePattern({
          id: 'update-test',
          patternText: 'updated',
          confidenceBase: 0.9,
          source: 'learned',
          hitCount: 3,
        }),
      )

      const patterns = await store.loadPatterns()
      const found = patterns.filter((p) => p.id === 'update-test')
      expect(found.length).toBe(1)
      expect(found[0]!.confidenceBase).toBe(0.9)
      expect(found[0]!.patternText).toBe('updated')
    })

    it('should not return disabled patterns', async () => {
      await store.savePattern(makePattern({ id: 'disabled-pat', enabled: false }))
      const patterns = await store.loadPatterns()
      const found = patterns.find((p) => p.id === 'disabled-pat')
      expect(found).toBeUndefined()
    })
  })

  describe('getStats()', () => {
    it('should return stats with zero counts on an empty store', async () => {
      const stats = await store.getStats()
      expect(stats).toBeDefined()
      expect(typeof stats.totalPatterns).toBe('number')
      expect(typeof stats.totalIncidents).toBe('number')
      expect(stats.totalPatterns).toBe(0)
      expect(stats.totalIncidents).toBe(0)
    })

    it('should reflect saved patterns in totalPatterns', async () => {
      await store.savePattern(makePattern({ id: 'stats-test-1' }))
      const stats = await store.getStats()
      expect(stats.totalPatterns).toBeGreaterThan(0)
    })

    it('should count patterns by source', async () => {
      await store.savePattern(makePattern({ id: 'builtin-1', source: 'builtin' }))
      await store.savePattern(makePattern({ id: 'learned-1', source: 'learned' }))
      const stats = await store.getStats()
      expect(stats.builtinPatterns).toBeGreaterThanOrEqual(1)
      expect(stats.learnedPatterns).toBeGreaterThanOrEqual(1)
    })

    it('should have a topPatterns array', async () => {
      const stats = await store.getStats()
      expect(Array.isArray(stats.topPatterns)).toBe(true)
    })
  })

  describe('store() — scan result ingestion', () => {
    it('should store a scan result without throwing', async () => {
      const result = makeScanResult({
        id: 'scan-001',
        input: 'ignore all previous instructions',
        detected: true,
        threatLevel: 'high',
        killChainPhase: 'initial_access',
        healingApplied: false,
      })
      await expect(store.store(result)).resolves.not.toThrow()
    })

    it('should store a false-negative candidate without throwing', async () => {
      const result = makeScanResult({
        id: 'scan-fn-001',
        input: 'How do I encode base64 in Python?',
        detected: false,
        threatLevel: 'none',
        killChainPhase: 'none',
        action: 'allow',
      })
      await expect(store.store(result)).resolves.not.toThrow()
    })

    it('should store multiple results without throwing', async () => {
      for (let i = 0; i < 10; i++) {
        await expect(store.store(makeScanResult({ id: `scan-multi-${i}` }))).resolves.not.toThrow()
      }
    })
  })

  describe('updateConfidence()', () => {
    it('should increase confidence by delta', async () => {
      await store.savePattern(makePattern({ id: 'conf-test', confidenceBase: 0.5 }))
      await store.updateConfidence('conf-test', 0.2)

      const patterns = await store.loadPatterns()
      const found = patterns.find((p) => p.id === 'conf-test')
      expect(found).toBeDefined()
      expect(found!.confidenceBase).toBeCloseTo(0.7, 5)
    })

    it('should clamp confidence to [0.1, 0.99] on large positive delta', async () => {
      await store.savePattern(makePattern({ id: 'clamp-high', confidenceBase: 0.95 }))
      await store.updateConfidence('clamp-high', 0.5)

      const patterns = await store.loadPatterns()
      const found = patterns.find((p) => p.id === 'clamp-high')
      expect(found!.confidenceBase).toBeLessThanOrEqual(0.99)
    })

    it('should clamp confidence to [0.1, 0.99] on large negative delta', async () => {
      await store.savePattern(makePattern({ id: 'clamp-low', confidenceBase: 0.15 }))
      await store.updateConfidence('clamp-low', -0.5)

      const patterns = await store.loadPatterns()
      const found = patterns.find((p) => p.id === 'clamp-low')
      expect(found!.confidenceBase).toBeGreaterThanOrEqual(0.1)
    })

    it('should be a no-op for unknown pattern id', async () => {
      await expect(store.updateConfidence('nonexistent-id', 0.1)).resolves.not.toThrow()
    })
  })

  describe('incrementHitCount()', () => {
    it('should increment hit count by 1', async () => {
      await store.savePattern(makePattern({ id: 'hit-test', hitCount: 3 }))
      await store.incrementHitCount('hit-test')

      const patterns = await store.loadPatterns()
      const found = patterns.find((p) => p.id === 'hit-test')
      expect(found!.hitCount).toBe(4)
    })

    it('should be a no-op for unknown pattern id', async () => {
      await expect(store.incrementHitCount('unknown-id')).resolves.not.toThrow()
    })
  })

  describe('incrementFalsePositiveCount()', () => {
    it('should increment false positive count by 1', async () => {
      await store.savePattern(makePattern({ id: 'fp-test', falsePositiveCount: 1 }))
      await store.incrementFalsePositiveCount('fp-test')

      const patterns = await store.loadPatterns()
      const found = patterns.find((p) => p.id === 'fp-test')
      expect(found!.falsePositiveCount).toBe(2)
    })
  })
})