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

package ucca

import (
	"testing"
)

func TestNIS2Module_Classification(t *testing.T) {
	module, err := NewNIS2Module()
	if err != nil {
		t.Fatalf("Failed to create NIS2 module: %v", err)
	}

	tests := []struct {
		name           string
		facts          *UnifiedFacts
		expectedClass  NIS2Classification
		expectedApply  bool
	}{
		{
			name: "Large energy company - Essential Entity",
			facts: &UnifiedFacts{
				Organization: OrganizationFacts{
					EmployeeCount: 500,
					AnnualRevenue: 100_000_000,
					Country:       "DE",
					EUMember:      true,
				},
				Sector: SectorFacts{
					PrimarySector: "energy",
				},
			},
			expectedClass: NIS2EssentialEntity,
			expectedApply: true,
		},
		{
			name: "Medium energy company - Important Entity",
			facts: &UnifiedFacts{
				Organization: OrganizationFacts{
					EmployeeCount: 100,
					AnnualRevenue: 20_000_000,
					Country:       "DE",
					EUMember:      true,
				},
				Sector: SectorFacts{
					PrimarySector: "energy",
				},
			},
			expectedClass: NIS2ImportantEntity,
			expectedApply: true,
		},
		{
			name: "Small energy company - Not Affected",
			facts: &UnifiedFacts{
				Organization: OrganizationFacts{
					EmployeeCount: 20,
					AnnualRevenue: 5_000_000,
					Country:       "DE",
					EUMember:      true,
				},
				Sector: SectorFacts{
					PrimarySector: "energy",
				},
			},
			expectedClass: NIS2NotAffected,
			expectedApply: false,
		},
		{
			name: "Large healthcare company - Essential Entity",
			facts: &UnifiedFacts{
				Organization: OrganizationFacts{
					EmployeeCount: 300,
					AnnualRevenue: 80_000_000,
					Country:       "DE",
					EUMember:      true,
				},
				Sector: SectorFacts{
					PrimarySector: "health",
				},
			},
			expectedClass: NIS2EssentialEntity,
			expectedApply: true,
		},
		{
			name: "Medium manufacturing company (Annex II) - Important Entity",
			facts: &UnifiedFacts{
				Organization: OrganizationFacts{
					EmployeeCount: 150,
					AnnualRevenue: 30_000_000,
					Country:       "DE",
					EUMember:      true,
				},
				Sector: SectorFacts{
					PrimarySector: "manufacturing",
				},
			},
			expectedClass: NIS2ImportantEntity,
			expectedApply: true,
		},
		{
			name: "Large manufacturing company (Annex II) - Important Entity (not Essential)",
			facts: &UnifiedFacts{
				Organization: OrganizationFacts{
					EmployeeCount: 500,
					AnnualRevenue: 100_000_000,
					Country:       "DE",
					EUMember:      true,
				},
				Sector: SectorFacts{
					PrimarySector: "manufacturing",
				},
			},
			expectedClass: NIS2ImportantEntity,
			expectedApply: true,
		},
		{
			name: "Retail company - Not Affected (not in NIS2 sectors)",
			facts: &UnifiedFacts{
				Organization: OrganizationFacts{
					EmployeeCount: 500,
					AnnualRevenue: 100_000_000,
					Country:       "DE",
					EUMember:      true,
				},
				Sector: SectorFacts{
					PrimarySector: "retail",
				},
			},
			expectedClass: NIS2NotAffected,
			expectedApply: false,
		},
		{
			name: "Small KRITIS operator - Essential Entity (exception)",
			facts: &UnifiedFacts{
				Organization: OrganizationFacts{
					EmployeeCount: 30,
					AnnualRevenue: 8_000_000,
					Country:       "DE",
					EUMember:      true,
				},
				Sector: SectorFacts{
					PrimarySector:      "energy",
					IsKRITIS:           true,
					KRITISThresholdMet: true,
				},
			},
			expectedClass: NIS2EssentialEntity,
			expectedApply: true,
		},
		{
			name: "Cloud provider (special service) - Essential Entity regardless of size",
			facts: &UnifiedFacts{
				Organization: OrganizationFacts{
					EmployeeCount: 20,
					AnnualRevenue: 5_000_000,
					Country:       "DE",
					EUMember:      true,
				},
				Sector: SectorFacts{
					PrimarySector:   "it_services",
					SpecialServices: []string{"cloud"},
				},
			},
			expectedClass: NIS2EssentialEntity,
			expectedApply: true,
		},
		{
			name: "DNS provider (special service) - Essential Entity regardless of size",
			facts: &UnifiedFacts{
				Organization: OrganizationFacts{
					EmployeeCount: 10,
					AnnualRevenue: 2_000_000,
					Country:       "DE",
					EUMember:      true,
				},
				Sector: SectorFacts{
					PrimarySector:   "digital_infrastructure",
					SpecialServices: []string{"dns"},
				},
			},
			expectedClass: NIS2EssentialEntity,
			expectedApply: true,
		},
		{
			name: "MSP provider - Essential Entity",
			facts: &UnifiedFacts{
				Organization: OrganizationFacts{
					EmployeeCount: 15,
					AnnualRevenue: 3_000_000,
					Country:       "DE",
					EUMember:      true,
				},
				Sector: SectorFacts{
					PrimarySector:   "ict_service_mgmt",
					SpecialServices: []string{"msp"},
				},
			},
			expectedClass: NIS2EssentialEntity,
			expectedApply: true,
		},
		{
			name: "Group company (counts group size) - Essential Entity",
			facts: &UnifiedFacts{
				Organization: OrganizationFacts{
					EmployeeCount:  50,
					AnnualRevenue:  15_000_000,
					IsPartOfGroup:  true,
					GroupEmployees: 500,
					GroupRevenue:   200_000_000,
					Country:        "DE",
					EUMember:       true,
				},
				Sector: SectorFacts{
					PrimarySector: "energy",
				},
			},
			expectedClass: NIS2EssentialEntity,
			expectedApply: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			classification := module.Classify(tt.facts)
			if classification != tt.expectedClass {
				t.Errorf("Classify() = %v, want %v", classification, tt.expectedClass)
			}

			isApplicable := module.IsApplicable(tt.facts)
			if isApplicable != tt.expectedApply {
				t.Errorf("IsApplicable() = %v, want %v", isApplicable, tt.expectedApply)
			}
		})
	}
}

func TestNIS2Module_DeriveObligations(t *testing.T) {
	module, err := NewNIS2Module()
	if err != nil {
		t.Fatalf("Failed to create NIS2 module: %v", err)
	}

	tests := []struct {
		name              string
		facts             *UnifiedFacts
		minObligations    int
		expectedContains  []string // Obligation IDs that should be present
		expectedMissing   []string // Obligation IDs that should NOT be present
	}{
		{
			name: "Essential Entity should have all obligations including audits",
			facts: &UnifiedFacts{
				Organization: OrganizationFacts{
					EmployeeCount: 500,
					AnnualRevenue: 100_000_000,
				},
				Sector: SectorFacts{
					PrimarySector: "energy",
				},
			},
			minObligations:   10,
			expectedContains: []string{"NIS2-OBL-001", "NIS2-OBL-002", "NIS2-OBL-012"}, // BSI registration, risk mgmt, audits
		},
		{
			name: "Important Entity should have obligations but NOT audit requirement",
			facts: &UnifiedFacts{
				Organization: OrganizationFacts{
					EmployeeCount: 100,
					AnnualRevenue: 20_000_000,
				},
				Sector: SectorFacts{
					PrimarySector: "manufacturing",
				},
			},
			minObligations:   8,
			expectedContains: []string{"NIS2-OBL-001", "NIS2-OBL-002"},
			expectedMissing:  []string{"NIS2-OBL-012"}, // Audits only for essential entities
		},
		{
			name: "Not affected entity should have no obligations",
			facts: &UnifiedFacts{
				Organization: OrganizationFacts{
					EmployeeCount: 20,
					AnnualRevenue: 5_000_000,
				},
				Sector: SectorFacts{
					PrimarySector: "retail",
				},
			},
			minObligations: 0,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			obligations := module.DeriveObligations(tt.facts)

			if len(obligations) < tt.minObligations {
				t.Errorf("DeriveObligations() returned %d obligations, want at least %d", len(obligations), tt.minObligations)
			}

			// Check expected contains
			for _, expectedID := range tt.expectedContains {
				found := false
				for _, obl := range obligations {
					if obl.ID == expectedID {
						found = true
						break
					}
				}
				if !found {
					t.Errorf("Expected obligation %s not found in results", expectedID)
				}
			}

			// Check expected missing
			for _, missingID := range tt.expectedMissing {
				for _, obl := range obligations {
					if obl.ID == missingID {
						t.Errorf("Obligation %s should NOT be present for this classification", missingID)
					}
				}
			}
		})
	}
}

func TestNIS2Module_IncidentDeadlines(t *testing.T) {
	module, err := NewNIS2Module()
	if err != nil {
		t.Fatalf("Failed to create NIS2 module: %v", err)
	}

	// Test for affected entity
	affectedFacts := &UnifiedFacts{
		Organization: OrganizationFacts{
			EmployeeCount: 100,
			AnnualRevenue: 20_000_000,
		},
		Sector: SectorFacts{
			PrimarySector: "energy",
		},
	}

	deadlines := module.GetIncidentDeadlines(affectedFacts)
	if len(deadlines) != 3 {
		t.Errorf("Expected 3 incident deadlines, got %d", len(deadlines))
	}

	// Check phases
	expectedPhases := map[string]string{
		"Frühwarnung":    "24 Stunden",
		"Vorfallmeldung": "72 Stunden",
		"Abschlussbericht": "1 Monat",
	}

	for _, deadline := range deadlines {
		if expected, ok := expectedPhases[deadline.Phase]; ok {
			if deadline.Deadline != expected {
				t.Errorf("Phase %s: expected deadline %s, got %s", deadline.Phase, expected, deadline.Deadline)
			}
		}
	}

	// Test for not affected entity
	notAffectedFacts := &UnifiedFacts{
		Organization: OrganizationFacts{
			EmployeeCount: 20,
			AnnualRevenue: 5_000_000,
		},
		Sector: SectorFacts{
			PrimarySector: "retail",
		},
	}

	deadlines = module.GetIncidentDeadlines(notAffectedFacts)
	if len(deadlines) != 0 {
		t.Errorf("Expected 0 incident deadlines for not affected entity, got %d", len(deadlines))
	}
}

func TestNIS2Module_DecisionTree(t *testing.T) {
	module, err := NewNIS2Module()
	if err != nil {
		t.Fatalf("Failed to create NIS2 module: %v", err)
	}

	tree := module.GetDecisionTree()
	if tree == nil {
		t.Fatal("Expected decision tree, got nil")
	}

	if tree.ID != "nis2_applicability" {
		t.Errorf("Expected tree ID 'nis2_applicability', got %s", tree.ID)
	}

	if tree.RootNode == nil {
		t.Fatal("Expected root node, got nil")
	}

	// Verify structure
	if tree.RootNode.YesNode == nil || tree.RootNode.NoNode == nil {
		t.Error("Root node should have both yes and no branches")
	}
}

func TestNIS2Module_DeriveControls(t *testing.T) {
	module, err := NewNIS2Module()
	if err != nil {
		t.Fatalf("Failed to create NIS2 module: %v", err)
	}

	// Test for affected entity
	affectedFacts := &UnifiedFacts{
		Organization: OrganizationFacts{
			EmployeeCount: 100,
			AnnualRevenue: 20_000_000,
		},
		Sector: SectorFacts{
			PrimarySector: "energy",
		},
	}

	controls := module.DeriveControls(affectedFacts)
	if len(controls) == 0 {
		t.Error("Expected controls for affected entity, got none")
	}

	// Check that controls have regulation ID set
	for _, ctrl := range controls {
		if ctrl.RegulationID != "nis2" {
			t.Errorf("Control %s has wrong regulation ID: %s", ctrl.ID, ctrl.RegulationID)
		}
	}

	// Test for not affected entity
	notAffectedFacts := &UnifiedFacts{
		Organization: OrganizationFacts{
			EmployeeCount: 20,
			AnnualRevenue: 5_000_000,
		},
		Sector: SectorFacts{
			PrimarySector: "retail",
		},
	}

	controls = module.DeriveControls(notAffectedFacts)
	if len(controls) != 0 {
		t.Errorf("Expected 0 controls for not affected entity, got %d", len(controls))
	}
}

func TestOrganizationFacts_SizeCalculations(t *testing.T) {
	tests := []struct {
		name            string
		org             OrganizationFacts
		expectedSize    string
		expectedSME     bool
		expectedNIS2    bool
		expectedLarge   bool
	}{
		{
			name: "Micro enterprise",
			org: OrganizationFacts{
				EmployeeCount: 5,
				AnnualRevenue: 1_000_000,
			},
			expectedSize:  "micro",
			expectedSME:   true,
			expectedNIS2:  false,
			expectedLarge: false,
		},
		{
			name: "Small enterprise",
			org: OrganizationFacts{
				EmployeeCount: 30,
				AnnualRevenue: 8_000_000,
			},
			expectedSize:  "small",
			expectedSME:   true,
			expectedNIS2:  false,
			expectedLarge: false,
		},
		{
			name: "Medium enterprise (by employees)",
			org: OrganizationFacts{
				EmployeeCount: 100,
				AnnualRevenue: 20_000_000,
			},
			expectedSize:  "medium",
			expectedSME:   true,
			expectedNIS2:  true,
			expectedLarge: false,
		},
		{
			name: "Medium enterprise (by revenue/balance)",
			org: OrganizationFacts{
				EmployeeCount:     40,
				AnnualRevenue:     15_000_000,
				BalanceSheetTotal: 15_000_000,
			},
			expectedSize:  "medium",    // < 50 employees BUT revenue AND balance > €10m → not small
			expectedSME:   true,
			expectedNIS2:  true,        // >= €10m revenue AND balance
			expectedLarge: false,
		},
		{
			name: "Large enterprise",
			org: OrganizationFacts{
				EmployeeCount: 500,
				AnnualRevenue: 100_000_000,
			},
			expectedSize:  "large",
			expectedSME:   false,
			expectedNIS2:  true,
			expectedLarge: true,
		},
		{
			name: "Group company (uses group figures)",
			org: OrganizationFacts{
				EmployeeCount:  50,
				AnnualRevenue:  15_000_000,
				IsPartOfGroup:  true,
				GroupEmployees: 500,
				GroupRevenue:   200_000_000,
			},
			expectedSize:  "large",
			expectedSME:   false,
			expectedNIS2:  true,
			expectedLarge: true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			size := tt.org.CalculateSizeCategory()
			if size != tt.expectedSize {
				t.Errorf("CalculateSizeCategory() = %v, want %v", size, tt.expectedSize)
			}

			isSME := tt.org.IsSME()
			if isSME != tt.expectedSME {
				t.Errorf("IsSME() = %v, want %v", isSME, tt.expectedSME)
			}

			meetsNIS2 := tt.org.MeetsNIS2SizeThreshold()
			if meetsNIS2 != tt.expectedNIS2 {
				t.Errorf("MeetsNIS2SizeThreshold() = %v, want %v", meetsNIS2, tt.expectedNIS2)
			}

			isLarge := tt.org.MeetsNIS2LargeThreshold()
			if isLarge != tt.expectedLarge {
				t.Errorf("MeetsNIS2LargeThreshold() = %v, want %v", isLarge, tt.expectedLarge)
			}
		})
	}
}