breakpilot-compliance/ai-compliance-sdk/internal/ucca/decision_tree_engine_test.go

package ucca

import (
	"testing"
)

func TestBuildDecisionTreeDefinition_ReturnsValidTree(t *testing.T) {
	tree := BuildDecisionTreeDefinition()

	if tree == nil {
		t.Fatal("Expected non-nil tree definition")
	}
	if tree.ID != "ai_act_two_axis" {
		t.Errorf("Expected ID 'ai_act_two_axis', got '%s'", tree.ID)
	}
	if tree.Version != "1.0.0" {
		t.Errorf("Expected version '1.0.0', got '%s'", tree.Version)
	}
	if len(tree.Questions) != 12 {
		t.Errorf("Expected 12 questions, got %d", len(tree.Questions))
	}

	// Check axis distribution
	hrCount := 0
	gpaiCount := 0
	for _, q := range tree.Questions {
		switch q.Axis {
		case "high_risk":
			hrCount++
		case "gpai":
			gpaiCount++
		default:
			t.Errorf("Unexpected axis '%s' for question %s", q.Axis, q.ID)
		}
	}
	if hrCount != 7 {
		t.Errorf("Expected 7 high_risk questions, got %d", hrCount)
	}
	if gpaiCount != 5 {
		t.Errorf("Expected 5 gpai questions, got %d", gpaiCount)
	}

	// Check all questions have required fields
	for _, q := range tree.Questions {
		if q.ID == "" {
			t.Error("Question has empty ID")
		}
		if q.Question == "" {
			t.Errorf("Question %s has empty question text", q.ID)
		}
		if q.Description == "" {
			t.Errorf("Question %s has empty description", q.ID)
		}
		if q.ArticleRef == "" {
			t.Errorf("Question %s has empty article_ref", q.ID)
		}
	}
}

func TestEvaluateDecisionTree_NotApplicable(t *testing.T) {
	// Q1=No → AI Act not applicable
	req := &DecisionTreeEvalRequest{
		SystemName: "Test System",
		Answers: map[string]DecisionTreeAnswer{
			Q1: {QuestionID: Q1, Value: false},
		},
	}

	result := EvaluateDecisionTree(req)

	if result.HighRiskResult != AIActNotApplicable {
		t.Errorf("Expected not_applicable, got %s", result.HighRiskResult)
	}
	if result.GPAIResult.IsGPAI {
		t.Error("Expected GPAI to be false when Q8 is not answered")
	}
	if result.SystemName != "Test System" {
		t.Errorf("Expected system name 'Test System', got '%s'", result.SystemName)
	}
}

func TestEvaluateDecisionTree_MinimalRisk(t *testing.T) {
	// Q1=Yes, Q2-Q7=No → minimal risk
	req := &DecisionTreeEvalRequest{
		SystemName: "Simple Tool",
		Answers: map[string]DecisionTreeAnswer{
			Q1: {QuestionID: Q1, Value: true},
			Q2: {QuestionID: Q2, Value: false},
			Q3: {QuestionID: Q3, Value: false},
			Q4: {QuestionID: Q4, Value: false},
			Q5: {QuestionID: Q5, Value: false},
			Q6: {QuestionID: Q6, Value: false},
			Q7: {QuestionID: Q7, Value: false},
			Q8: {QuestionID: Q8, Value: false},
		},
	}

	result := EvaluateDecisionTree(req)

	if result.HighRiskResult != AIActMinimalRisk {
		t.Errorf("Expected minimal_risk, got %s", result.HighRiskResult)
	}
	if result.GPAIResult.IsGPAI {
		t.Error("Expected GPAI to be false")
	}
	if result.GPAIResult.Category != GPAICategoryNone {
		t.Errorf("Expected GPAI category 'none', got '%s'", result.GPAIResult.Category)
	}
}

func TestEvaluateDecisionTree_HighRisk_Biometric(t *testing.T) {
	// Q1=Yes, Q2=Yes → high risk (biometric)
	req := &DecisionTreeEvalRequest{
		SystemName: "Face Recognition",
		Answers: map[string]DecisionTreeAnswer{
			Q1: {QuestionID: Q1, Value: true},
			Q2: {QuestionID: Q2, Value: true},
			Q3: {QuestionID: Q3, Value: false},
			Q4: {QuestionID: Q4, Value: false},
			Q5: {QuestionID: Q5, Value: false},
			Q6: {QuestionID: Q6, Value: false},
			Q7: {QuestionID: Q7, Value: false},
		},
	}

	result := EvaluateDecisionTree(req)

	if result.HighRiskResult != AIActHighRisk {
		t.Errorf("Expected high_risk, got %s", result.HighRiskResult)
	}

	// Should have high-risk obligations
	if len(result.CombinedObligations) == 0 {
		t.Error("Expected non-empty obligations for high-risk system")
	}
}

func TestEvaluateDecisionTree_HighRisk_CriticalInfrastructure(t *testing.T) {
	// Q1=Yes, Q3=Yes → high risk (critical infrastructure)
	req := &DecisionTreeEvalRequest{
		SystemName: "Energy Grid AI",
		Answers: map[string]DecisionTreeAnswer{
			Q1: {QuestionID: Q1, Value: true},
			Q2: {QuestionID: Q2, Value: false},
			Q3: {QuestionID: Q3, Value: true},
			Q4: {QuestionID: Q4, Value: false},
			Q5: {QuestionID: Q5, Value: false},
			Q6: {QuestionID: Q6, Value: false},
			Q7: {QuestionID: Q7, Value: false},
		},
	}

	result := EvaluateDecisionTree(req)

	if result.HighRiskResult != AIActHighRisk {
		t.Errorf("Expected high_risk, got %s", result.HighRiskResult)
	}
}

func TestEvaluateDecisionTree_HighRisk_Education(t *testing.T) {
	// Q1=Yes, Q4=Yes → high risk (education/employment)
	req := &DecisionTreeEvalRequest{
		SystemName: "Exam Grading AI",
		Answers: map[string]DecisionTreeAnswer{
			Q1: {QuestionID: Q1, Value: true},
			Q2: {QuestionID: Q2, Value: false},
			Q3: {QuestionID: Q3, Value: false},
			Q4: {QuestionID: Q4, Value: true},
		},
	}

	result := EvaluateDecisionTree(req)

	if result.HighRiskResult != AIActHighRisk {
		t.Errorf("Expected high_risk, got %s", result.HighRiskResult)
	}
}

func TestEvaluateDecisionTree_HighRisk_AutonomousDecisions(t *testing.T) {
	// Q1=Yes, Q7=Yes → high risk (autonomous decisions)
	req := &DecisionTreeEvalRequest{
		SystemName: "Credit Scoring AI",
		Answers: map[string]DecisionTreeAnswer{
			Q1: {QuestionID: Q1, Value: true},
			Q2: {QuestionID: Q2, Value: false},
			Q3: {QuestionID: Q3, Value: false},
			Q4: {QuestionID: Q4, Value: false},
			Q5: {QuestionID: Q5, Value: false},
			Q6: {QuestionID: Q6, Value: false},
			Q7: {QuestionID: Q7, Value: true},
		},
	}

	result := EvaluateDecisionTree(req)

	if result.HighRiskResult != AIActHighRisk {
		t.Errorf("Expected high_risk, got %s", result.HighRiskResult)
	}
}

func TestEvaluateDecisionTree_GPAI_Standard(t *testing.T) {
	// Q8=Yes, Q10=No → GPAI standard
	req := &DecisionTreeEvalRequest{
		SystemName: "Custom LLM",
		Answers: map[string]DecisionTreeAnswer{
			Q1:  {QuestionID: Q1, Value: true},
			Q8:  {QuestionID: Q8, Value: true},
			Q9:  {QuestionID: Q9, Value: true},
			Q10: {QuestionID: Q10, Value: false},
			Q11: {QuestionID: Q11, Value: false},
			Q12: {QuestionID: Q12, Value: false},
		},
	}

	result := EvaluateDecisionTree(req)

	if !result.GPAIResult.IsGPAI {
		t.Error("Expected IsGPAI to be true")
	}
	if result.GPAIResult.Category != GPAICategoryStandard {
		t.Errorf("Expected category 'standard', got '%s'", result.GPAIResult.Category)
	}
	if result.GPAIResult.IsSystemicRisk {
		t.Error("Expected IsSystemicRisk to be false")
	}

	// Should have Art. 51, 53, 50 (generative)
	hasArt51 := false
	hasArt53 := false
	hasArt50 := false
	for _, a := range result.GPAIResult.ApplicableArticles {
		if a == "Art. 51" {
			hasArt51 = true
		}
		if a == "Art. 53" {
			hasArt53 = true
		}
		if a == "Art. 50" {
			hasArt50 = true
		}
	}
	if !hasArt51 {
		t.Error("Expected Art. 51 in applicable articles")
	}
	if !hasArt53 {
		t.Error("Expected Art. 53 in applicable articles")
	}
	if !hasArt50 {
		t.Error("Expected Art. 50 in applicable articles (generative AI)")
	}
}

func TestEvaluateDecisionTree_GPAI_SystemicRisk(t *testing.T) {
	// Q8=Yes, Q10=Yes → GPAI systemic risk
	req := &DecisionTreeEvalRequest{
		SystemName: "GPT-5",
		Answers: map[string]DecisionTreeAnswer{
			Q1:  {QuestionID: Q1, Value: true},
			Q8:  {QuestionID: Q8, Value: true},
			Q9:  {QuestionID: Q9, Value: true},
			Q10: {QuestionID: Q10, Value: true},
			Q11: {QuestionID: Q11, Value: true},
			Q12: {QuestionID: Q12, Value: true},
		},
	}

	result := EvaluateDecisionTree(req)

	if !result.GPAIResult.IsGPAI {
		t.Error("Expected IsGPAI to be true")
	}
	if result.GPAIResult.Category != GPAICategorySystemic {
		t.Errorf("Expected category 'systemic', got '%s'", result.GPAIResult.Category)
	}
	if !result.GPAIResult.IsSystemicRisk {
		t.Error("Expected IsSystemicRisk to be true")
	}

	// Should have Art. 55
	hasArt55 := false
	for _, a := range result.GPAIResult.ApplicableArticles {
		if a == "Art. 55" {
			hasArt55 = true
		}
	}
	if !hasArt55 {
		t.Error("Expected Art. 55 in applicable articles (systemic risk)")
	}
}

func TestEvaluateDecisionTree_Combined_HighRiskAndGPAI(t *testing.T) {
	// Q1=Yes, Q4=Yes (high risk) + Q8=Yes, Q9=Yes (GPAI standard)
	req := &DecisionTreeEvalRequest{
		SystemName:        "HR Screening with LLM",
		SystemDescription: "LLM-based applicant screening system",
		Answers: map[string]DecisionTreeAnswer{
			Q1:  {QuestionID: Q1, Value: true},
			Q2:  {QuestionID: Q2, Value: false},
			Q3:  {QuestionID: Q3, Value: false},
			Q4:  {QuestionID: Q4, Value: true},
			Q5:  {QuestionID: Q5, Value: false},
			Q6:  {QuestionID: Q6, Value: false},
			Q7:  {QuestionID: Q7, Value: true},
			Q8:  {QuestionID: Q8, Value: true},
			Q9:  {QuestionID: Q9, Value: true},
			Q10: {QuestionID: Q10, Value: false},
			Q11: {QuestionID: Q11, Value: false},
			Q12: {QuestionID: Q12, Value: false},
		},
	}

	result := EvaluateDecisionTree(req)

	// Both axes should be triggered
	if result.HighRiskResult != AIActHighRisk {
		t.Errorf("Expected high_risk, got %s", result.HighRiskResult)
	}
	if !result.GPAIResult.IsGPAI {
		t.Error("Expected GPAI to be true")
	}
	if result.GPAIResult.Category != GPAICategoryStandard {
		t.Errorf("Expected GPAI category 'standard', got '%s'", result.GPAIResult.Category)
	}

	// Combined obligations should include both axes
	if len(result.CombinedObligations) < 5 {
		t.Errorf("Expected at least 5 combined obligations, got %d", len(result.CombinedObligations))
	}

	// Should have articles from both axes
	if len(result.ApplicableArticles) < 3 {
		t.Errorf("Expected at least 3 applicable articles, got %d", len(result.ApplicableArticles))
	}

	// Check system name preserved
	if result.SystemName != "HR Screening with LLM" {
		t.Errorf("Expected system name preserved, got '%s'", result.SystemName)
	}
	if result.SystemDescription != "LLM-based applicant screening system" {
		t.Errorf("Expected description preserved, got '%s'", result.SystemDescription)
	}
}

func TestEvaluateDecisionTree_GPAI_MarketPenetration(t *testing.T) {
	// Q8=Yes, Q10=No, Q12=Yes → GPAI standard with market penetration warning
	req := &DecisionTreeEvalRequest{
		SystemName: "Popular Chatbot",
		Answers: map[string]DecisionTreeAnswer{
			Q1:  {QuestionID: Q1, Value: true},
			Q8:  {QuestionID: Q8, Value: true},
			Q9:  {QuestionID: Q9, Value: true},
			Q10: {QuestionID: Q10, Value: false},
			Q11: {QuestionID: Q11, Value: true},
			Q12: {QuestionID: Q12, Value: true},
		},
	}

	result := EvaluateDecisionTree(req)

	if result.GPAIResult.Category != GPAICategoryStandard {
		t.Errorf("Expected category 'standard' (not systemic because Q10=No), got '%s'", result.GPAIResult.Category)
	}

	// Should have Art. 51 Abs. 3 warning
	hasArt51_3 := false
	for _, a := range result.GPAIResult.ApplicableArticles {
		if a == "Art. 51 Abs. 3" {
			hasArt51_3 = true
		}
	}
	if !hasArt51_3 {
		t.Error("Expected Art. 51 Abs. 3 in applicable articles for high market penetration")
	}
}

func TestEvaluateDecisionTree_NoGPAI(t *testing.T) {
	// Q8=No → No GPAI classification
	req := &DecisionTreeEvalRequest{
		SystemName: "Traditional ML",
		Answers: map[string]DecisionTreeAnswer{
			Q1: {QuestionID: Q1, Value: true},
			Q8: {QuestionID: Q8, Value: false},
		},
	}

	result := EvaluateDecisionTree(req)

	if result.GPAIResult.IsGPAI {
		t.Error("Expected IsGPAI to be false")
	}
	if result.GPAIResult.Category != GPAICategoryNone {
		t.Errorf("Expected category 'none', got '%s'", result.GPAIResult.Category)
	}
	if len(result.GPAIResult.Obligations) != 0 {
		t.Errorf("Expected 0 GPAI obligations, got %d", len(result.GPAIResult.Obligations))
	}
}

func TestAnswerIsYes(t *testing.T) {
	tests := []struct {
		name     string
		answers  map[string]DecisionTreeAnswer
		qID      string
		expected bool
	}{
		{"yes answer", map[string]DecisionTreeAnswer{"Q1": {Value: true}}, "Q1", true},
		{"no answer", map[string]DecisionTreeAnswer{"Q1": {Value: false}}, "Q1", false},
		{"missing answer", map[string]DecisionTreeAnswer{}, "Q1", false},
		{"different question", map[string]DecisionTreeAnswer{"Q2": {Value: true}}, "Q1", false},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := answerIsYes(tt.answers, tt.qID)
			if result != tt.expected {
				t.Errorf("Expected %v, got %v", tt.expected, result)
			}
		})
	}
}