fetch_ml/tests/benchmarks/kms_benchmark_test.go

package benchmarks

import (
	"testing"
	"time"

	"github.com/jfraeys/fetch_ml/internal/crypto"
	"github.com/jfraeys/fetch_ml/internal/crypto/kms"
	kmsconfig "github.com/jfraeys/fetch_ml/internal/crypto/kms/config"
)

// BenchmarkEncryptArtifact measures the full encryption pipeline performance.
// Per ADR-012: Total overhead should be <10ms for MemoryProvider.
func BenchmarkEncryptArtifact(b *testing.B) {
	tkm := crypto.NewTestTenantKeyManager(nil)

	// Provision a test tenant
	hierarchy, err := tkm.ProvisionTenant("bench-tenant")
	if err != nil {
		b.Fatalf("Failed to provision tenant: %v", err)
	}

	// Test data - 1KB payload (typical model weights chunk)
	plaintext := make([]byte, 1024)
	for i := range plaintext {
		plaintext[i] = byte(i % 256)
	}

	b.ResetTimer()
	b.ReportAllocs()

	for b.Loop() {
		_, err := tkm.EncryptArtifact("bench-tenant", "artifact-1", hierarchy.KMSKeyID, plaintext)
		if err != nil {
			b.Fatalf("Encrypt failed: %v", err)
		}
	}
}

// BenchmarkDecryptArtifact measures the full decryption pipeline performance.
// Per ADR-012: Total overhead should be <10ms for MemoryProvider.
func BenchmarkDecryptArtifact(b *testing.B) {
	tkm := crypto.NewTestTenantKeyManager(nil)

	hierarchy, err := tkm.ProvisionTenant("bench-tenant")
	if err != nil {
		b.Fatalf("Failed to provision tenant: %v", err)
	}

	plaintext := make([]byte, 1024)
	for i := range plaintext {
		plaintext[i] = byte(i % 256)
	}

	// Pre-encrypt data
	encrypted, err := tkm.EncryptArtifact("bench-tenant", "artifact-1", hierarchy.KMSKeyID, plaintext)
	if err != nil {
		b.Fatalf("Pre-encryption failed: %v", err)
	}

	b.ReportAllocs()

	for b.Loop() {
		_, err := tkm.DecryptArtifact(encrypted, hierarchy.KMSKeyID)
		if err != nil {
			b.Fatalf("Decrypt failed: %v", err)
		}
	}
}

// BenchmarkMemoryProvider_Encrypt measures baseline encryption without network overhead.
// This establishes the theoretical minimum for KMS operations.
func BenchmarkMemoryProvider_Encrypt(b *testing.B) {
	provider := kms.NewMemoryProvider()
	defer provider.Close()

	cache := kms.NewDEKCache(kmsconfig.DefaultCacheConfig())
	defer cache.Clear()

	config := kmsconfig.Config{
		Provider: kmsconfig.ProviderTypeMemory,
		Cache:    kmsconfig.DefaultCacheConfig(),
	}

	tkm := crypto.NewTenantKeyManager(provider, cache, config, nil)

	hierarchy, err := tkm.ProvisionTenant("bench-tenant")
	if err != nil {
		b.Fatalf("Failed to provision tenant: %v", err)
	}

	plaintext := make([]byte, 1024)

	b.ReportAllocs()

	for b.Loop() {
		_, err := tkm.EncryptArtifact("bench-tenant", "artifact-1", hierarchy.KMSKeyID, plaintext)
		if err != nil {
			b.Fatalf("Encrypt failed: %v", err)
		}
	}
}

// BenchmarkCacheHit verifies cached DEKs provide <10ms overhead.
func BenchmarkCacheHit(b *testing.B) {
	tkm := crypto.NewTestTenantKeyManager(nil)

	hierarchy, err := tkm.ProvisionTenant("bench-tenant")
	if err != nil {
		b.Fatalf("Failed to provision tenant: %v", err)
	}

	plaintext := make([]byte, 1024)

	// First encrypt to populate cache
	encrypted, err := tkm.EncryptArtifact("bench-tenant", "cached-artifact", hierarchy.KMSKeyID, plaintext)
	if err != nil {
		b.Fatalf("Pre-encryption failed: %v", err)
	}

	// First decrypt to populate DEK cache
	_, err = tkm.DecryptArtifact(encrypted, hierarchy.KMSKeyID)
	if err != nil {
		b.Fatalf("First decrypt failed: %v", err)
	}

	b.ReportAllocs()

	for b.Loop() {
		_, err := tkm.DecryptArtifact(encrypted, hierarchy.KMSKeyID)
		if err != nil {
			b.Fatalf("Decrypt failed: %v", err)
		}
	}
}

// BenchmarkKeyRotation measures key rotation overhead.
func BenchmarkKeyRotation(b *testing.B) {
	tkm := crypto.NewTestTenantKeyManager(nil)

	hierarchy, err := tkm.ProvisionTenant("bench-tenant")
	if err != nil {
		b.Fatalf("Failed to provision tenant: %v", err)
	}

	b.ReportAllocs()

	for b.Loop() {
		// Rotate key
		newHierarchy, err := tkm.RotateTenantKey("bench-tenant", hierarchy)
		if err != nil {
			b.Fatalf("Rotation failed: %v", err)
		}
		hierarchy = newHierarchy
	}
}

// BenchmarkEncryptArtifact_LargePayload measures encryption with larger payloads.
func BenchmarkEncryptArtifact_LargePayload(b *testing.B) {
	tkm := crypto.NewTestTenantKeyManager(nil)

	hierarchy, err := tkm.ProvisionTenant("bench-tenant")
	if err != nil {
		b.Fatalf("Failed to provision tenant: %v", err)
	}

	// 1MB payload
	plaintext := make([]byte, 1024*1024)

	b.ReportAllocs()

	for b.Loop() {
		_, err := tkm.EncryptArtifact("bench-tenant", "large-artifact", hierarchy.KMSKeyID, plaintext)
		if err != nil {
			b.Fatalf("Encrypt failed: %v", err)
		}
	}
}

// BenchmarkParallelEncrypt measures concurrent encryption performance.
func BenchmarkParallelEncrypt(b *testing.B) {
	tkm := crypto.NewTestTenantKeyManager(nil)

	hierarchy, err := tkm.ProvisionTenant("bench-tenant")
	if err != nil {
		b.Fatalf("Failed to provision tenant: %v", err)
	}

	plaintext := make([]byte, 1024)

	b.ResetTimer()
	b.ReportAllocs()

	b.RunParallel(func(pb *testing.PB) {
		i := 0
		for pb.Next() {
			artifactID := "parallel-artifact-" + string(rune('0'+i%10))
			_, err := tkm.EncryptArtifact("bench-tenant", artifactID, hierarchy.KMSKeyID, plaintext)
			if err != nil {
				b.Fatalf("Encrypt failed: %v", err)
			}
			i++
		}
	})
}

// VerifyPerformanceRequirement runs a quick sanity check for the <10ms requirement.
// This is not a benchmark but a verification that typical operations complete within limits.
func TestEncryptPerformance_10msRequirement(t *testing.T) {
	tkm := crypto.NewTestTenantKeyManager(nil)

	hierarchy, err := tkm.ProvisionTenant("perf-test-tenant")
	if err != nil {
		t.Fatalf("Failed to provision tenant: %v", err)
	}

	plaintext := make([]byte, 1024)

	// Warm up
	for range 10 {
		_, _ = tkm.EncryptArtifact("perf-test-tenant", "warmup", hierarchy.KMSKeyID, plaintext)
	}

	// Measure 100 operations
	start := time.Now()
	for i := 0; i < 100; i++ {
		_, err := tkm.EncryptArtifact("perf-test-tenant", "perf-test", hierarchy.KMSKeyID, plaintext)
		if err != nil {
			t.Fatalf("Encrypt failed: %v", err)
		}
	}
	elapsed := time.Since(start)

	avgPerOp := elapsed / 100
	if avgPerOp > 10*time.Millisecond {
		t.Errorf("Average encrypt time %v exceeds 10ms requirement", avgPerOp)
	}

	t.Logf("Average encrypt time: %v (requirement: <10ms)", avgPerOp)
}

// TestDecryptPerformance_10msRequirement verifies decrypt completes within 10ms.
func TestDecryptPerformance_10msRequirement(t *testing.T) {
	tkm := crypto.NewTestTenantKeyManager(nil)

	hierarchy, err := tkm.ProvisionTenant("perf-test-tenant")
	if err != nil {
		t.Fatalf("Failed to provision tenant: %v", err)
	}

	plaintext := make([]byte, 1024)

	// Pre-encrypt
	encrypted, err := tkm.EncryptArtifact("perf-test-tenant", "perf-test", hierarchy.KMSKeyID, plaintext)
	if err != nil {
		t.Fatalf("Pre-encryption failed: %v", err)
	}

	// Warm up cache
	for range 10 {
		_, _ = tkm.DecryptArtifact(encrypted, hierarchy.KMSKeyID)
	}

	// Measure 100 operations with cache
	start := time.Now()
	for range 10 {
		_, err := tkm.DecryptArtifact(encrypted, hierarchy.KMSKeyID)
		if err != nil {
			t.Fatalf("Decrypt failed: %v", err)
		}
	}
	elapsed := time.Since(start)

	avgPerOp := elapsed / 100
	if avgPerOp > 10*time.Millisecond {
		t.Errorf("Average decrypt time %v exceeds 10ms requirement", avgPerOp)
	}

	t.Logf("Average decrypt time: %v (requirement: <10ms)", avgPerOp)
}