rony-llm-agent/pkg/rag/backends/chroma/chroma_test.go

package chroma_test

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

	"github.com/VictorVargas/rony-llm-agent/pkg/rag/backends/chroma"
)

func TestBackend_Upsert(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/api/v1/collections/rony-memory/upsert" {
			w.WriteHeader(http.StatusOK)
			w.Write([]byte(`{}`))
			return
		}
		w.WriteHeader(http.StatusNotFound)
	}))
	defer server.Close()

	backend, err := chroma.New(chroma.Config{BaseURL: server.URL})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	err = backend.Upsert(context.Background(), "test-id", []float32{0.1, 0.2}, map[string]string{"key": "value"})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestBackend_Upsert_APIError(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte(`{"error":"server error"}`))
	}))
	defer server.Close()

	backend, err := chroma.New(chroma.Config{BaseURL: server.URL})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	err = backend.Upsert(context.Background(), "test-id", []float32{0.1}, nil)
	if err == nil {
		t.Fatal("expected error")
	}
}

func TestBackend_Search(t *testing.T) {
	meta := []map[string]interface{}{{"key": "value"}}
	metaNested := [][]map[string]interface{}{meta}
	mockResponse := map[string]interface{}{
		"names":   []string{"rony-memory"},
		"results": []map[string]interface{}{
			{
				"ids":       [][]string{{"test-id"}},
				"documents": [][]string{{"test content"}},
				"distances": [][]float64{{0.9}},
				"metadatas": metaNested,
			},
		},
	}

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/api/v1/collections/rony-memory/query" {
			w.Header().Set("Content-Type", "application/json")
			json.NewEncoder(w).Encode(mockResponse)
			return
		}
		w.WriteHeader(http.StatusNotFound)
	}))
	defer server.Close()

	backend, err := chroma.New(chroma.Config{BaseURL: server.URL})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	results, err := backend.Search(context.Background(), []float32{0.1, 0.2}, 5)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

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

	if results[0].ID != "test-id" {
		t.Errorf("expected 'test-id', got %q", results[0].ID)
	}

	if results[0].Content != "test content" {
		t.Errorf("expected 'test content', got %q", results[0].Content)
	}
}

func TestBackend_ForgetAll(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/api/v1/collections/rony-memory/delete" {
			w.WriteHeader(http.StatusOK)
			w.Write([]byte(`{}`))
			return
		}
		w.WriteHeader(http.StatusNotFound)
	}))
	defer server.Close()

	backend, err := chroma.New(chroma.Config{BaseURL: server.URL})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	err = backend.ForgetAll(context.Background())
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestBackend_ForgetAll_APIError(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusInternalServerError)
		w.Write([]byte(`{"error":"server error"}`))
	}))
	defer server.Close()

	backend, err := chroma.New(chroma.Config{BaseURL: server.URL})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	err = backend.ForgetAll(context.Background())
	if err == nil {
		t.Fatal("expected error")
	}
}

func TestNew(t *testing.T) {
	_, err := chroma.New(chroma.Config{BaseURL: ""})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	_, err = chroma.New(chroma.Config{})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestNew_WithTimeout(t *testing.T) {
	backend, err := chroma.New(chroma.Config{
		BaseURL: "http://localhost:8000",
		Timeout: 10,
	})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if backend == nil {
		t.Fatal("expected non-nil backend")
	}
}

func TestUpsert_InvalidJSON(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
		w.Write([]byte(`{}`))
	}))
	defer server.Close()

	backend, err := chroma.New(chroma.Config{BaseURL: server.URL})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	// Test with nil metadata (should work)
	err = backend.Upsert(context.Background(), "test-id", []float32{0.1}, nil)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}
}

func TestSearch_EmptyResults(t *testing.T) {
	mockResponse := map[string]interface{}{
		"results": []map[string]interface{}{},
	}

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/api/v1/collections/rony-memory/query" {
			w.Header().Set("Content-Type", "application/json")
			json.NewEncoder(w).Encode(mockResponse)
			return
		}
		w.WriteHeader(http.StatusNotFound)
	}))
	defer server.Close()

	backend, err := chroma.New(chroma.Config{BaseURL: server.URL})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	results, err := backend.Search(context.Background(), []float32{0.1, 0.2}, 5)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if len(results) != 0 {
		t.Fatalf("expected 0 results, got %d", len(results))
	}
}

func TestSearch_MalformedResponse(t *testing.T) {
	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/api/v1/collections/rony-memory/query" {
			w.Header().Set("Content-Type", "application/json")
			w.Write([]byte(`invalid json`))
			return
		}
		w.WriteHeader(http.StatusNotFound)
	}))
	defer server.Close()

	backend, err := chroma.New(chroma.Config{BaseURL: server.URL})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	_, err = backend.Search(context.Background(), []float32{0.1, 0.2}, 5)
	if err == nil {
		t.Fatal("expected error for malformed response")
	}
}

func TestSearch_MissingFields(t *testing.T) {
	mockResponse := map[string]interface{}{
		"results": []map[string]interface{}{},
	}

	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path == "/api/v1/collections/rony-memory/query" {
			w.Header().Set("Content-Type", "application/json")
			json.NewEncoder(w).Encode(mockResponse)
			return
		}
		w.WriteHeader(http.StatusNotFound)
	}))
	defer server.Close()

	backend, err := chroma.New(chroma.Config{BaseURL: server.URL})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	results, err := backend.Search(context.Background(), []float32{0.1, 0.2}, 5)
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	if len(results) != 0 {
		t.Fatalf("expected 0 results, got %d", len(results))
	}
}

func TestUpsert_ContextCanceled(t *testing.T) {
	backend, err := chroma.New(chroma.Config{BaseURL: "http://localhost:8000"})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	err = backend.Upsert(ctx, "test-id", []float32{0.1}, nil)
	if err == nil {
		t.Fatal("expected error for canceled context")
	}
}

func TestSearch_ContextCanceled(t *testing.T) {
	backend, err := chroma.New(chroma.Config{BaseURL: "http://localhost:8000"})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	_, err = backend.Search(ctx, []float32{0.1, 0.2}, 5)
	if err == nil {
		t.Fatal("expected error for canceled context")
	}
}

func TestForget_ContextCanceled(t *testing.T) {
	backend, err := chroma.New(chroma.Config{BaseURL: "http://localhost:8000"})
	if err != nil {
		t.Fatalf("unexpected error: %v", err)
	}

	ctx, cancel := context.WithCancel(context.Background())
	cancel()

	err = backend.Forget(ctx, "test-id")
	if err == nil {
		t.Fatal("expected error for canceled context")
	}
}