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

package ucca

import (
	"os"
	"testing"
)

func TestRouterBaseCollections(t *testing.T) {
	c := &LegalRAGClient{}
	os.Unsetenv("RAG_ROUTER_COLLECTIONS")
	def := c.routerBaseCollections()
	if len(def) != 6 || def[1] != "bp_compliance_ce" {
		t.Fatalf("default base collections unexpected: %v", def)
	}

	os.Setenv("RAG_ROUTER_COLLECTIONS", " bp_compliance_ce , kb_2026_1_build ,, ")
	defer os.Unsetenv("RAG_ROUTER_COLLECTIONS")
	got := c.routerBaseCollections()
	if len(got) != 2 || got[0] != "bp_compliance_ce" || got[1] != "kb_2026_1_build" {
		t.Fatalf("env override parse failed (trim/empty): %v", got)
	}
}

func TestRouterSliceSelection(t *testing.T) {
	// The router appends the slice exactly when the query is in scope (inKBScope) and routing is on.
	// Mirror the selection logic so a regression in either is caught without a live Qdrant.
	c := &LegalRAGClient{kbSliceCollection: "kb_2026_1_build", kbScopeRoutingEnabled: true}
	sel := func(q string) bool {
		colls := c.routerBaseCollections()
		if c.kbScopeRoutingEnabled && c.kbSliceCollection != "" && inKBScope(q) {
			colls = append(colls, c.kbSliceCollection)
		}
		for _, x := range colls {
			if x == c.kbSliceCollection {
				return true
			}
		}
		return false
	}
	if !sel("Welche neun Kriterien nennt WP248 fuer ein hohes Risiko?") {
		t.Error("in-scope guidance query must include the slice")
	}
	if sel("Was sagt NIST SP 800-53 zu Access Control?") {
		t.Error("out-of-scope query must NOT include the slice")
	}
	c.kbScopeRoutingEnabled = false
	if sel("Welche Kriterien nennt WP248?") {
		t.Error("routing disabled => slice never included")
	}
}

func TestBalanceByRegulation(t *testing.T) {
	regs := []detectedRegulation{
		{Canonical: "CRA", CodeValues: []string{"CRA"}},
		{Canonical: "MaschVO", CodeValues: []string{"MASCHVO", "MVO", "MACHINERY"}},
	}
	// CRA dominates by score; without balancing the top-4 would be all CRA + NIST.
	pool := []LegalSearchResult{
		{RegulationCode: "CRA", Score: 0.99},
		{RegulationCode: "CRA", Score: 0.98},
		{RegulationCode: "CRA", Score: 0.97},
		{RegulationCode: "NIST", Score: 0.96},
		{RegulationCode: "MACHINERY", Score: 0.70},
		{RegulationCode: "MVO", Score: 0.65},
	}
	out := balanceByRegulation(pool, regs, 4)
	var hasCRA, hasMasch bool
	for _, r := range out {
		switch r.RegulationCode {
		case "CRA":
			hasCRA = true
		case "MACHINERY", "MVO":
			hasMasch = true
		}
	}
	if !hasCRA || !hasMasch {
		t.Errorf("both named domains must be represented: CRA=%v MaschVO=%v out=%v", hasCRA, hasMasch, out)
	}
	if out[0].RegulationCode != "CRA" || !(out[1].RegulationCode == "MACHINERY" || out[1].RegulationCode == "MVO") {
		t.Errorf("round-robin should alternate domains, got %s then %s", out[0].RegulationCode, out[1].RegulationCode)
	}
}

func TestDedupResults(t *testing.T) {
	in := []LegalSearchResult{
		{RegulationCode: "EDPB WP248", ArticleLabel: "III.B", Text: "lorem", Score: 0.7},
		{RegulationCode: "EDPB WP248", ArticleLabel: "III.B", Text: "lorem", Score: 0.9}, // dup, higher score
		{RegulationCode: "DSGVO", ArticleLabel: "Art. 35", Text: "ipsum", Score: 0.8},
	}
	out := dedupResults(in)
	if len(out) != 2 {
		t.Fatalf("expected 2 deduped, got %d", len(out))
	}
	for _, r := range out {
		if r.RegulationCode == "EDPB WP248" && r.Score != 0.9 {
			t.Errorf("dedup must keep highest score, got %v", r.Score)
		}
	}
}