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

package ucca

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
)

func TestSearchCollection_UsesCorrectCollection(t *testing.T) {
	// Track which collection was requested
	var requestedURL string

	// Mock Ollama (embedding)
	ollamaMock := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		json.NewEncoder(w).Encode(ollamaEmbeddingResponse{
			Embedding: make([]float64, 1024),
		})
	}))
	defer ollamaMock.Close()

	// Mock Qdrant
	qdrantMock := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requestedURL = r.URL.Path
		json.NewEncoder(w).Encode(qdrantSearchResponse{
			Result: []qdrantSearchHit{},
		})
	}))
	defer qdrantMock.Close()

	// Parse qdrant mock host/port
	qdrantAddr := strings.TrimPrefix(qdrantMock.URL, "http://")
	parts := strings.Split(qdrantAddr, ":")

	client := &LegalRAGClient{
		qdrantHost:     parts[0],
		qdrantPort:     parts[1],
		ollamaURL:      ollamaMock.URL,
		embeddingModel: "bge-m3",
		collection:     "bp_compliance_ce",
		httpClient:     http.DefaultClient,
	}

	// Test with explicit collection
	_, err := client.SearchCollection(context.Background(), "bp_compliance_recht", "test query", nil, 3)
	if err != nil {
		t.Fatalf("SearchCollection failed: %v", err)
	}

	if !strings.Contains(requestedURL, "/collections/bp_compliance_recht/") {
		t.Errorf("Expected collection bp_compliance_recht in URL, got: %s", requestedURL)
	}
}

func TestSearchCollection_FallbackDefault(t *testing.T) {
	var requestedURL string

	ollamaMock := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		json.NewEncoder(w).Encode(ollamaEmbeddingResponse{
			Embedding: make([]float64, 1024),
		})
	}))
	defer ollamaMock.Close()

	qdrantMock := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requestedURL = r.URL.Path
		json.NewEncoder(w).Encode(qdrantSearchResponse{
			Result: []qdrantSearchHit{},
		})
	}))
	defer qdrantMock.Close()

	qdrantAddr := strings.TrimPrefix(qdrantMock.URL, "http://")
	parts := strings.Split(qdrantAddr, ":")

	client := &LegalRAGClient{
		qdrantHost:     parts[0],
		qdrantPort:     parts[1],
		ollamaURL:      ollamaMock.URL,
		embeddingModel: "bge-m3",
		collection:     "bp_compliance_ce",
		httpClient:     http.DefaultClient,
	}

	// Test with empty collection (should fall back to default)
	_, err := client.SearchCollection(context.Background(), "", "test query", nil, 3)
	if err != nil {
		t.Fatalf("SearchCollection failed: %v", err)
	}

	if !strings.Contains(requestedURL, "/collections/bp_compliance_ce/") {
		t.Errorf("Expected default collection bp_compliance_ce in URL, got: %s", requestedURL)
	}
}

func TestSearch_StillWorks(t *testing.T) {
	var requestedURL string

	ollamaMock := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		json.NewEncoder(w).Encode(ollamaEmbeddingResponse{
			Embedding: make([]float64, 1024),
		})
	}))
	defer ollamaMock.Close()

	qdrantMock := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		requestedURL = r.URL.Path
		json.NewEncoder(w).Encode(qdrantSearchResponse{
			Result: []qdrantSearchHit{
				{
					ID:    "1",
					Score: 0.95,
					Payload: map[string]interface{}{
						"chunk_text":         "Test content",
						"regulation_id":      "eu_2016_679",
						"regulation_name_de": "DSGVO",
						"regulation_short":   "DSGVO",
						"category":           "regulation",
						"source":             "https://example.com",
					},
				},
			},
		})
	}))
	defer qdrantMock.Close()

	qdrantAddr := strings.TrimPrefix(qdrantMock.URL, "http://")
	parts := strings.Split(qdrantAddr, ":")

	client := &LegalRAGClient{
		qdrantHost:     parts[0],
		qdrantPort:     parts[1],
		ollamaURL:      ollamaMock.URL,
		embeddingModel: "bge-m3",
		collection:     "bp_compliance_ce",
		httpClient:     http.DefaultClient,
	}

	results, err := client.Search(context.Background(), "DSGVO Art. 35", nil, 5)
	if err != nil {
		t.Fatalf("Search failed: %v", err)
	}

	if len(results) != 1 {
		t.Fatalf("Expected 1 result, got %d", len(results))
	}

	if results[0].RegulationCode != "eu_2016_679" {
		t.Errorf("Expected regulation_code eu_2016_679, got %s", results[0].RegulationCode)
	}

	if !strings.Contains(requestedURL, "/collections/bp_compliance_ce/") {
		t.Errorf("Expected default collection in URL, got: %s", requestedURL)
	}
}