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

package ucca

import "testing"

func TestClassifyRole(t *testing.T) {
	tests := []struct {
		name string
		r    LegalSearchResult
		want string
	}{
		{"NIST -> control_standard", LegalSearchResult{RegulationShort: "NIST SP 800-82r3", ArticleLabel: "AU-8"}, roleControlStandard},
		{"OWASP -> control_standard", LegalSearchResult{RegulationShort: "OWASP ASVS"}, roleControlStandard},
		{"CRA Anhang -> operational_requirement", LegalSearchResult{RegulationShort: "CRA", ArticleLabel: "CRA Anhang I", Category: "regulation"}, roleOperationalReq},
		{"CRA Meldepflicht -> procedural_requirement", LegalSearchResult{RegulationShort: "CRA", ArticleLabel: "Art. 14 CRA Meldepflicht", Category: "regulation"}, roleProceduralReq},
		{"ENISA Good Practices -> implementation_guidance", LegalSearchResult{RegulationShort: "ENISA Supply Chain Good Practices"}, roleImplGuidance},
		{"EDPB Leitlinie -> interpretation", LegalSearchResult{RegulationShort: "EDPB DPO", ArticleLabel: "WP243 Leitlinien Datenschutzbeauftragte"}, roleInterpretation},
		{"DORA article -> obligation", LegalSearchResult{RegulationShort: "DORA", ArticleLabel: "Art. 5 DORA", Category: "regulation"}, roleObligation},
		{"DSGVO Begriffsbestimmungen -> definition", LegalSearchResult{RegulationShort: "DSGVO", ArticleLabel: "Art. 4 DSGVO Begriffsbestimmungen", Category: "regulation"}, roleDefinition},
		{"recital -> definition", LegalSearchResult{RegulationShort: "CRA", IsRecital: true}, roleDefinition},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			if got := classifyRole(tt.r); got != tt.want {
				t.Errorf("classifyRole() = %q, want %q", got, tt.want)
			}
		})
	}
}

func TestApplyControlRoles_PoolPreference(t *testing.T) {
	// op_req > procedural > control_standard > impl_guidance; non-control roles get no boost.
	roles := []struct {
		r        LegalSearchResult
		wantGain float64
	}{
		{LegalSearchResult{ArticleLabel: "CRA Anhang I", Category: "regulation"}, controlPoolGain + 0.100},
		{LegalSearchResult{ArticleLabel: "Art. 14 CRA Meldepflicht", Category: "regulation"}, controlPoolGain + 0.075},
		{LegalSearchResult{RegulationShort: "NIST SP 800-53"}, controlPoolGain + 0.050},
		{LegalSearchResult{RegulationShort: "ENISA Good Practices"}, controlPoolGain + 0.000},
		{LegalSearchResult{ArticleLabel: "Art. 5 DORA", Category: "regulation"}, 0.0}, // obligation: no boost
	}
	for _, rc := range roles {
		out := []LegalSearchResult{rc.r}
		out[0].Score = 1.0
		applyControlRoles(out)
		if got := out[0].Score - 1.0; got < rc.wantGain-1e-9 || got > rc.wantGain+1e-9 {
			t.Errorf("role %q: gain %.3f, want %.3f", classifyRole(rc.r), got, rc.wantGain)
		}
	}
}

func TestIsControlPoolRole(t *testing.T) {
	for _, r := range []string{roleOperationalReq, roleProceduralReq, roleControlStandard, roleImplGuidance} {
		if !isControlPoolRole(r) {
			t.Errorf("%q should be in the control-pool", r)
		}
	}
	for _, r := range []string{roleObligation, roleInterpretation, roleDefinition} {
		if isControlPoolRole(r) {
			t.Errorf("%q should NOT be in the control-pool", r)
		}
	}
}

func TestControlRoleOf_Payload(t *testing.T) {
	// searchControls filters its deep dense pull by classifying the raw Qdrant payload.
	nist := map[string]interface{}{"regulation_short": "NIST SP 800-82r3", "article": "AU-8"}
	if got := controlRoleOf(nist); got != roleControlStandard {
		t.Errorf("untagged NIST payload role = %q, want control_standard", got)
	}
	craAnnex := map[string]interface{}{"regulation_short": "CRA", "article": "Anhang-I", "category": "regulation"}
	if got := controlRoleOf(craAnnex); got != roleOperationalReq {
		t.Errorf("CRA Anhang payload role = %q, want operational_requirement", got)
	}
	dora := map[string]interface{}{"regulation_short": "DORA", "article_label": "Art. 5 DORA", "category": "regulation"}
	if got := controlRoleOf(dora); isControlPoolRole(got) {
		t.Errorf("DORA abstract article role = %q must be excluded from the control-pool", got)
	}
}

func headHasRole(head []LegalSearchResult, role string) bool {
	for _, r := range head {
		if classifyRole(r) == role {
			return true
		}
	}
	return false
}

func TestEnsureControlDiversity(t *testing.T) {
	ig := func(n string) LegalSearchResult {
		return LegalSearchResult{RegulationShort: "ENISA " + n + " Good Practices"}
	}
	opReq := LegalSearchResult{RegulationShort: "CRA", ArticleLabel: "CRA Anhang I", Category: "regulation"}
	std := LegalSearchResult{RegulationShort: "NIST SP 800-53"}

	t.Run("injects missing op_req + control_standard, guidance keeps Top-1", func(t *testing.T) {
		out := ensureControlDiversity([]LegalSearchResult{ig("A"), ig("B"), ig("C"), std, opReq}, 3)
		head := out[:3]
		if classifyRole(head[0]) != roleImplGuidance {
			t.Errorf("Top-1 should stay implementation_guidance, got %q", classifyRole(head[0]))
		}
		if !headHasRole(head, roleOperationalReq) {
			t.Error("top-K must contain an operational_requirement after diversity")
		}
		if !headHasRole(head, roleControlStandard) {
			t.Error("top-K must contain a control_standard after diversity")
		}
	})

	t.Run("no-op when both roles already present", func(t *testing.T) {
		out := ensureControlDiversity([]LegalSearchResult{opReq, std, ig("A"), ig("B")}, 3)
		if classifyRole(out[0]) != roleOperationalReq || classifyRole(out[1]) != roleControlStandard {
			t.Error("already-diverse top-K must be left untouched")
		}
	})

	t.Run("absent role is not forced (no panic)", func(t *testing.T) {
		out := ensureControlDiversity([]LegalSearchResult{ig("A"), ig("B"), ig("C"), std}, 3)
		if !headHasRole(out[:3], roleControlStandard) {
			t.Error("present control_standard should be injected")
		}
		if headHasRole(out[:3], roleOperationalReq) {
			t.Error("operational_requirement absent from the pool must NOT appear")
		}
	})

	t.Run("topK covering the whole pool is unchanged", func(t *testing.T) {
		out := ensureControlDiversity([]LegalSearchResult{ig("A"), opReq}, 5)
		if len(out) != 2 || classifyRole(out[0]) != roleImplGuidance {
			t.Error("topK >= len must return results unchanged")
		}
	})
}