breakpilot-compliance/pca-platform/heuristic-service/internal/stepup/pow_test.go

package stepup

import (
	"testing"

	"github.com/breakpilot/pca-platform/heuristic-service/internal/config"
)

func TestNewPoWService(t *testing.T) {
	cfg := &config.PoWConfig{
		Enabled:       true,
		Difficulty:    4,
		MaxDurationMs: 5000,
	}

	service := NewPoWService(cfg)

	if service == nil {
		t.Fatal("Expected non-nil service")
	}
	if !service.IsEnabled() {
		t.Error("Expected service to be enabled")
	}
	if service.GetDifficulty() != 4 {
		t.Errorf("Expected difficulty 4, got %d", service.GetDifficulty())
	}
}

func TestCreateChallenge(t *testing.T) {
	cfg := &config.PoWConfig{
		Enabled:       true,
		Difficulty:    4,
		MaxDurationMs: 5000,
	}

	service := NewPoWService(cfg)
	response, err := service.CreateChallenge("test-session")

	if err != nil {
		t.Fatalf("Unexpected error: %v", err)
	}
	if response == nil {
		t.Fatal("Expected non-nil response")
	}
	if response.Challenge == "" {
		t.Error("Expected non-empty challenge")
	}
	if response.ChallengeID == "" {
		t.Error("Expected non-empty challenge ID")
	}
	if response.Difficulty != 4 {
		t.Errorf("Expected difficulty 4, got %d", response.Difficulty)
	}
	if response.MaxTimeMs != 5000 {
		t.Errorf("Expected max time 5000, got %d", response.MaxTimeMs)
	}
}

func TestVerifyProof_Valid(t *testing.T) {
	cfg := &config.PoWConfig{
		Enabled:       true,
		Difficulty:    2, // Low difficulty for fast testing
		MaxDurationMs: 5000,
	}

	service := NewPoWService(cfg)

	// Find a valid nonce for a known challenge
	challenge := "test-challenge-123"
	var validNonce int64 = -1

	// Brute force to find valid nonce (with low difficulty)
	for nonce := int64(0); nonce < 10000; nonce++ {
		if service.VerifyProof(challenge, nonce, 2) {
			validNonce = nonce
			break
		}
	}

	if validNonce == -1 {
		t.Skip("Could not find valid nonce in reasonable time")
	}

	// Verify the found nonce
	if !service.VerifyProof(challenge, validNonce, 2) {
		t.Errorf("Expected valid proof for nonce %d", validNonce)
	}
}

func TestVerifyProof_Invalid(t *testing.T) {
	cfg := &config.PoWConfig{
		Enabled:       true,
		Difficulty:    4,
		MaxDurationMs: 5000,
	}

	service := NewPoWService(cfg)

	// Nonce 0 is very unlikely to be valid for difficulty 4
	valid := service.VerifyProof("random-challenge", 0, 4)

	if valid {
		t.Error("Expected invalid proof for nonce 0")
	}
}

func TestVerifyChallenge_ValidFlow(t *testing.T) {
	cfg := &config.PoWConfig{
		Enabled:       true,
		Difficulty:    2,
		MaxDurationMs: 10000,
	}

	service := NewPoWService(cfg)

	// Create challenge
	response, err := service.CreateChallenge("test-session")
	if err != nil {
		t.Fatalf("Failed to create challenge: %v", err)
	}

	// Find valid nonce
	var validNonce int64 = -1
	for nonce := int64(0); nonce < 100000; nonce++ {
		if service.VerifyProof(response.Challenge, nonce, 2) {
			validNonce = nonce
			break
		}
	}

	if validNonce == -1 {
		t.Skip("Could not find valid nonce")
	}

	// Verify challenge
	req := &PoWVerifyRequest{
		SessionID:   "test-session",
		ChallengeID: response.ChallengeID,
		Challenge:   response.Challenge,
		Nonce:       validNonce,
	}

	verified, err := service.VerifyChallenge(req)
	if err != nil {
		t.Fatalf("Verification error: %v", err)
	}
	if !verified {
		t.Error("Expected verification to succeed")
	}
}

func TestVerifyChallenge_WrongSession(t *testing.T) {
	cfg := &config.PoWConfig{
		Enabled:       true,
		Difficulty:    2,
		MaxDurationMs: 5000,
	}

	service := NewPoWService(cfg)

	// Create challenge for session A
	response, _ := service.CreateChallenge("session-a")

	// Try to verify with session B
	req := &PoWVerifyRequest{
		SessionID:   "session-b",
		ChallengeID: response.ChallengeID,
		Challenge:   response.Challenge,
		Nonce:       0,
	}

	verified, _ := service.VerifyChallenge(req)
	if verified {
		t.Error("Expected verification to fail for wrong session")
	}
}

func TestVerifyChallenge_NonexistentChallenge(t *testing.T) {
	cfg := &config.PoWConfig{
		Enabled:       true,
		Difficulty:    2,
		MaxDurationMs: 5000,
	}

	service := NewPoWService(cfg)

	req := &PoWVerifyRequest{
		SessionID:   "test-session",
		ChallengeID: "nonexistent-challenge",
		Challenge:   "test",
		Nonce:       0,
	}

	verified, _ := service.VerifyChallenge(req)
	if verified {
		t.Error("Expected verification to fail for nonexistent challenge")
	}
}

func TestCleanupExpiredChallenges(t *testing.T) {
	cfg := &config.PoWConfig{
		Enabled:       true,
		Difficulty:    2,
		MaxDurationMs: 1, // Very short for testing
	}

	service := NewPoWService(cfg)

	// Create challenge
	service.CreateChallenge("test-session")

	if len(service.challenges) != 1 {
		t.Errorf("Expected 1 challenge, got %d", len(service.challenges))
	}

	// Wait for expiration
	// Note: In real test, we'd mock time or set ExpiresAt in the past

	// For now, just verify cleanup doesn't crash
	service.CleanupExpiredChallenges()
}

func TestIsEnabled(t *testing.T) {
	cfg := &config.PoWConfig{
		Enabled:       false,
		Difficulty:    4,
		MaxDurationMs: 5000,
	}

	service := NewPoWService(cfg)

	if service.IsEnabled() {
		t.Error("Expected service to be disabled")
	}
}