container-census/internal/vulnerability/cache_test.go

package vulnerability

import (
	"fmt"
	"testing"
	"time"
)

// mockStorage implements VulnerabilityStorage for testing
type mockStorage struct {
	scans map[string]*VulnerabilityScan
}

func newMockStorage() *mockStorage {
	return &mockStorage{
		scans: make(map[string]*VulnerabilityScan),
	}
}

func (m *mockStorage) GetVulnerabilityScan(imageID string) (*VulnerabilityScan, error) {
	scan, exists := m.scans[imageID]
	if !exists {
		return nil, fmt.Errorf("scan not found")
	}
	return scan, nil
}

func (m *mockStorage) SaveVulnerabilityScan(scan *VulnerabilityScan, vulnerabilities []Vulnerability) error {
	m.scans[scan.ImageID] = scan
	return nil
}

func TestCache_NeedsScan_FreshScan(t *testing.T) {
	storage := newMockStorage()
	cache := NewCache(storage, 24*time.Hour)

	// Save a fresh scan
	scan := &VulnerabilityScan{
		ImageID:   "sha256:abc123",
		ImageName: "nginx:latest",
		ScannedAt: time.Now(),
		Success:   true,
	}
	err := cache.Set(scan, []Vulnerability{})
	if err != nil {
		t.Fatalf("Failed to save scan: %v", err)
	}

	// Should NOT need scanning (fresh scan within TTL)
	if cache.NeedsScan("sha256:abc123") {
		t.Error("Expected NeedsScan to return false for fresh scan")
	}
}

func TestCache_NeedsScan_StaleScan(t *testing.T) {
	storage := newMockStorage()
	cache := NewCache(storage, 1*time.Hour) // 1 hour TTL

	// Save a stale scan (2 hours old)
	scan := &VulnerabilityScan{
		ImageID:   "sha256:abc123",
		ImageName: "nginx:latest",
		ScannedAt: time.Now().Add(-2 * time.Hour),
		Success:   true,
	}
	err := cache.Set(scan, []Vulnerability{})
	if err != nil {
		t.Fatalf("Failed to save scan: %v", err)
	}

	// Should need scanning (stale scan beyond TTL)
	if !cache.NeedsScan("sha256:abc123") {
		t.Error("Expected NeedsScan to return true for stale scan")
	}
}

func TestCache_NeedsScan_NoScan(t *testing.T) {
	storage := newMockStorage()
	cache := NewCache(storage, 24*time.Hour)

	// No scan exists for this image
	if !cache.NeedsScan("sha256:nonexistent") {
		t.Error("Expected NeedsScan to return true for non-existent scan")
	}
}

func TestCache_Get_FromMemoryCache(t *testing.T) {
	storage := newMockStorage()
	cache := NewCache(storage, 24*time.Hour)

	// Save a scan
	scan := &VulnerabilityScan{
		ImageID:   "sha256:abc123",
		ImageName: "nginx:latest",
		ScannedAt: time.Now(),
		Success:   true,
	}
	err := cache.Set(scan, []Vulnerability{})
	if err != nil {
		t.Fatalf("Failed to save scan: %v", err)
	}

	// Should retrieve from memory cache
	retrieved, err := cache.Get("sha256:abc123")
	if err != nil {
		t.Fatalf("Failed to get scan: %v", err)
	}

	if retrieved.ImageID != scan.ImageID {
		t.Errorf("Expected image ID %s, got %s", scan.ImageID, retrieved.ImageID)
	}

	// Verify it's using memory cache (storage should have been called once during Set)
	if cache.Size() != 1 {
		t.Errorf("Expected cache size 1, got %d", cache.Size())
	}
}

func TestCache_Get_FromDatabaseAfterMemoryCacheExpiry(t *testing.T) {
	storage := newMockStorage()
	cache := NewCache(storage, 24*time.Hour)

	// Save a scan
	scan := &VulnerabilityScan{
		ImageID:   "sha256:abc123",
		ImageName: "nginx:latest",
		ScannedAt: time.Now(),
		Success:   true,
	}
	err := cache.Set(scan, []Vulnerability{})
	if err != nil {
		t.Fatalf("Failed to save scan: %v", err)
	}

	// Clear memory cache
	cache.Clear()

	if cache.Size() != 0 {
		t.Errorf("Expected cache size 0 after clear, got %d", cache.Size())
	}

	// Should retrieve from database and re-populate memory cache
	retrieved, err := cache.Get("sha256:abc123")
	if err != nil {
		t.Fatalf("Failed to get scan from database: %v", err)
	}

	if retrieved.ImageID != scan.ImageID {
		t.Errorf("Expected image ID %s, got %s", scan.ImageID, retrieved.ImageID)
	}

	// Memory cache should be re-populated
	if cache.Size() != 1 {
		t.Errorf("Expected cache size 1 after database retrieval, got %d", cache.Size())
	}
}

func TestCache_Invalidate(t *testing.T) {
	storage := newMockStorage()
	cache := NewCache(storage, 24*time.Hour)

	// Save a scan
	scan := &VulnerabilityScan{
		ImageID:   "sha256:abc123",
		ImageName: "nginx:latest",
		ScannedAt: time.Now(),
		Success:   true,
	}
	err := cache.Set(scan, []Vulnerability{})
	if err != nil {
		t.Fatalf("Failed to save scan: %v", err)
	}

	// Invalidate the cache
	cache.Invalidate("sha256:abc123")

	// Should retrieve from database again
	retrieved, err := cache.Get("sha256:abc123")
	if err != nil {
		t.Fatalf("Failed to get scan after invalidation: %v", err)
	}

	if retrieved.ImageID != scan.ImageID {
		t.Errorf("Expected image ID %s, got %s", scan.ImageID, retrieved.ImageID)
	}
}

func TestCache_PruneExpired(t *testing.T) {
	storage := newMockStorage()
	cache := NewCache(storage, 1*time.Hour)

	// Save fresh scan
	freshScan := &VulnerabilityScan{
		ImageID:   "sha256:fresh",
		ImageName: "nginx:latest",
		ScannedAt: time.Now(),
		Success:   true,
	}
	err := cache.Set(freshScan, []Vulnerability{})
	if err != nil {
		t.Fatalf("Failed to save fresh scan: %v", err)
	}

	// Save stale scan (2 hours old)
	staleScan := &VulnerabilityScan{
		ImageID:   "sha256:stale",
		ImageName: "redis:latest",
		ScannedAt: time.Now().Add(-2 * time.Hour),
		Success:   true,
	}
	err = cache.Set(staleScan, []Vulnerability{})
	if err != nil {
		t.Fatalf("Failed to save stale scan: %v", err)
	}

	if cache.Size() != 2 {
		t.Errorf("Expected cache size 2, got %d", cache.Size())
	}

	// Prune expired entries
	removed := cache.PruneExpired()

	if removed != 1 {
		t.Errorf("Expected 1 entry to be pruned, got %d", removed)
	}

	if cache.Size() != 1 {
		t.Errorf("Expected cache size 1 after pruning, got %d", cache.Size())
	}

	// Fresh scan should still be in cache
	_, err = cache.Get("sha256:fresh")
	if err != nil {
		t.Errorf("Fresh scan should still be in cache: %v", err)
	}
}

func TestCache_TTLUpdate(t *testing.T) {
	storage := newMockStorage()
	cache := NewCache(storage, 24*time.Hour)

	if cache.GetTTL() != 24*time.Hour {
		t.Errorf("Expected TTL 24h, got %v", cache.GetTTL())
	}

	cache.SetTTL(48 * time.Hour)

	if cache.GetTTL() != 48*time.Hour {
		t.Errorf("Expected TTL 48h after update, got %v", cache.GetTTL())
	}
}

func TestCache_IsValid(t *testing.T) {
	storage := newMockStorage()
	cache := NewCache(storage, 1*time.Hour)

	// Fresh scan
	freshScan := &VulnerabilityScan{
		ImageID:   "sha256:fresh",
		ImageName: "nginx:latest",
		ScannedAt: time.Now(),
		Success:   true,
	}
	err := cache.Set(freshScan, []Vulnerability{})
	if err != nil {
		t.Fatalf("Failed to save fresh scan: %v", err)
	}

	// Stale scan
	staleScan := &VulnerabilityScan{
		ImageID:   "sha256:stale",
		ImageName: "redis:latest",
		ScannedAt: time.Now().Add(-2 * time.Hour),
		Success:   true,
	}
	err = cache.Set(staleScan, []Vulnerability{})
	if err != nil {
		t.Fatalf("Failed to save stale scan: %v", err)
	}

	// Fresh scan should be valid
	if !cache.IsValid("sha256:fresh") {
		t.Error("Expected fresh scan to be valid")
	}

	// Stale scan should be invalid
	if cache.IsValid("sha256:stale") {
		t.Error("Expected stale scan to be invalid")
	}

	// Non-existent scan should be invalid
	if cache.IsValid("sha256:nonexistent") {
		t.Error("Expected non-existent scan to be invalid")
	}
}

func TestCache_ConcurrentAccess(t *testing.T) {
	storage := newMockStorage()
	cache := NewCache(storage, 24*time.Hour)

	// Save initial scan
	scan := &VulnerabilityScan{
		ImageID:   "sha256:abc123",
		ImageName: "nginx:latest",
		ScannedAt: time.Now(),
		Success:   true,
	}
	err := cache.Set(scan, []Vulnerability{})
	if err != nil {
		t.Fatalf("Failed to save scan: %v", err)
	}

	// Concurrent reads
	done := make(chan bool)
	for i := 0; i < 10; i++ {
		go func() {
			_, err := cache.Get("sha256:abc123")
			if err != nil {
				t.Errorf("Concurrent read failed: %v", err)
			}
			done <- true
		}()
	}

	// Wait for all goroutines
	for i := 0; i < 10; i++ {
		<-done
	}
}