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}, "test content", 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}, "test content", 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(), "query", []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}, "test content", 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(), "query", []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(), "query", []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(), "query", []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}, "test content", 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, "query", []float32{0.1, 0.2}, 5) if err == nil { t.Fatal("expected error for canceled context") } } func TestSearch_NoVectorReturnsError(t *testing.T) { backend, err := chroma.New(chroma.Config{BaseURL: "http://localhost:8000"}) if err != nil { t.Fatalf("unexpected error: %v", err) } _, err = backend.Search(context.Background(), "query", nil, 5) if err == nil { t.Fatal("expected error since chroma has no lexical fallback") } } 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") } }