package db import ( "encoding/json" "fmt" "net/http" "net/http/httptest" "net/url" "os" "strings" "testing" "GoNavi-Wails/internal/connection" ) func newMockChromaServer(t *testing.T, handler http.HandlerFunc) *httptest.Server { t.Helper() server := httptest.NewServer(handler) t.Cleanup(server.Close) return server } func newTestChromaDB(t *testing.T, serverURL string) *ChromaDB { t.Helper() parsed, err := url.Parse(serverURL) if err != nil { t.Fatalf("parse server URL: %v", err) } host, port, ok := parseHostPortWithDefault(parsed.Host, defaultChromaPort) if !ok { t.Fatalf("parse host port failed: %s", parsed.Host) } db := &ChromaDB{} if err := db.Connect(connection.ConnectionConfig{ Type: "chroma", Host: host, Port: port, }); err != nil { t.Fatalf("connect chroma: %v", err) } t.Cleanup(func() { _ = db.Close() }) return db } func writeChromaJSON(w http.ResponseWriter, value interface{}) { w.Header().Set("Content-Type", "application/json") _ = json.NewEncoder(w).Encode(value) } func TestChromaConnectDetectsV2(t *testing.T) { server := newMockChromaServer(t, func(w http.ResponseWriter, r *http.Request) { if r.Method == http.MethodGet && r.URL.Path == "/api/v2/heartbeat" { writeChromaJSON(w, map[string]interface{}{"nanosecond heartbeat": 1}) return } w.WriteHeader(http.StatusNotFound) }) db := newTestChromaDB(t, server.URL) if db.apiVersion != 2 { t.Fatalf("apiVersion = %d, want 2", db.apiVersion) } } func TestChromaConnectFallsBackToV1(t *testing.T) { server := newMockChromaServer(t, func(w http.ResponseWriter, r *http.Request) { if r.URL.Path == "/api/v2/heartbeat" { w.WriteHeader(http.StatusNotFound) return } if r.Method == http.MethodGet && r.URL.Path == "/api/v1/heartbeat" { writeChromaJSON(w, map[string]interface{}{"nanosecond heartbeat": 1}) return } w.WriteHeader(http.StatusNotFound) }) db := newTestChromaDB(t, server.URL) if db.apiVersion != 1 { t.Fatalf("apiVersion = %d, want 1", db.apiVersion) } } func TestChromaGetDatabasesAndTablesV2(t *testing.T) { server := newMockChromaServer(t, func(w http.ResponseWriter, r *http.Request) { switch r.URL.Path { case "/api/v2/heartbeat": writeChromaJSON(w, map[string]interface{}{"ok": true}) case "/api/v2/tenants/default_tenant/databases": writeChromaJSON(w, []map[string]interface{}{ {"name": "analytics"}, {"name": "default_database"}, }) case "/api/v2/tenants/default_tenant/databases/default_database/collections": writeChromaJSON(w, []chromaCollection{ {ID: "col-products", Name: "products", Database: "default_database", Tenant: "default_tenant"}, {ID: "col-logs", Name: "logs", Database: "default_database", Tenant: "default_tenant"}, }) default: w.WriteHeader(http.StatusNotFound) } }) db := newTestChromaDB(t, server.URL) dbs, err := db.GetDatabases() if err != nil { t.Fatalf("GetDatabases failed: %v", err) } if strings.Join(dbs, ",") != "analytics,default_database" { t.Fatalf("databases = %v", dbs) } tables, err := db.GetTables("") if err != nil { t.Fatalf("GetTables failed: %v", err) } if strings.Join(tables, ",") != "logs,products" { t.Fatalf("tables = %v", tables) } } func TestChromaSelectConvertsToGetRows(t *testing.T) { var capturedPath string var capturedBody map[string]interface{} server := newMockChromaServer(t, func(w http.ResponseWriter, r *http.Request) { switch { case r.URL.Path == "/api/v2/heartbeat": writeChromaJSON(w, map[string]interface{}{"ok": true}) case r.URL.Path == "/api/v2/tenants/default_tenant/databases/default_database/collections": writeChromaJSON(w, []chromaCollection{{ID: "col-products", Name: "products", Database: "default_database"}}) case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/collections/col-products/get"): capturedPath = r.URL.Path _ = json.NewDecoder(r.Body).Decode(&capturedBody) writeChromaJSON(w, chromaGetResponse{ IDs: []string{"p1"}, Documents: []interface{}{"first product"}, Metadatas: []map[string]interface{}{{"category": "book", "price": json.Number("19.5")}}, }) default: w.WriteHeader(http.StatusNotFound) } }) db := newTestChromaDB(t, server.URL) rows, columns, err := db.Query(`SELECT * FROM "products" LIMIT 10 OFFSET 5`) if err != nil { t.Fatalf("Query failed: %v", err) } if capturedPath == "" { t.Fatal("expected get endpoint to be called") } if intFromAny(capturedBody["limit"], 0) != 10 || intFromAny(capturedBody["offset"], -1) != 5 { t.Fatalf("captured body = %#v", capturedBody) } if len(rows) != 1 || rows[0]["id"] != "p1" || rows[0]["metadata.category"] != "book" { t.Fatalf("rows = %#v", rows) } if !containsString(columns, "metadata.category") { t.Fatalf("columns missing metadata.category: %v", columns) } } func TestChromaJSONQueryFlattensResults(t *testing.T) { server := newMockChromaServer(t, func(w http.ResponseWriter, r *http.Request) { switch { case r.URL.Path == "/api/v2/heartbeat": writeChromaJSON(w, map[string]interface{}{"ok": true}) case r.URL.Path == "/api/v2/tenants/default_tenant/databases/default_database/collections": writeChromaJSON(w, []chromaCollection{{ID: "col-products", Name: "products", Database: "default_database"}}) case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/collections/col-products/query"): writeChromaJSON(w, map[string]interface{}{ "ids": [][]string{{"p1", "p2"}}, "documents": [][]string{{"first", "second"}}, "distances": [][]float64{{0.1, 0.2}}, "metadatas": [][]map[string]interface{}{{{"category": "book"}, {"category": "tool"}}}, }) default: w.WriteHeader(http.StatusNotFound) } }) db := newTestChromaDB(t, server.URL) rows, columns, err := db.Query(`{"query":"products","query_embeddings":[[0.1,0.2]],"n_results":2}`) if err != nil { t.Fatalf("Query failed: %v", err) } if len(rows) != 2 || rows[1]["id"] != "p2" || rows[1]["distance"] == nil { t.Fatalf("rows = %#v", rows) } if !containsString(columns, "distance") || !containsString(columns, "metadata.category") { t.Fatalf("columns = %v", columns) } } func TestChromaApplyChangesUpsertAndDelete(t *testing.T) { var upsertBody map[string]interface{} var deleteBody map[string]interface{} server := newMockChromaServer(t, func(w http.ResponseWriter, r *http.Request) { switch { case r.URL.Path == "/api/v2/heartbeat": writeChromaJSON(w, map[string]interface{}{"ok": true}) case r.URL.Path == "/api/v2/tenants/default_tenant/databases/default_database/collections": writeChromaJSON(w, []chromaCollection{{ID: "col-products", Name: "products", Database: "default_database"}}) case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/collections/col-products/upsert"): _ = json.NewDecoder(r.Body).Decode(&upsertBody) writeChromaJSON(w, map[string]interface{}{"ok": true}) case r.Method == http.MethodPost && strings.HasSuffix(r.URL.Path, "/collections/col-products/delete"): _ = json.NewDecoder(r.Body).Decode(&deleteBody) writeChromaJSON(w, map[string]interface{}{"ok": true}) default: w.WriteHeader(http.StatusNotFound) } }) db := newTestChromaDB(t, server.URL) err := db.ApplyChanges("products", connection.ChangeSet{ Deletes: []map[string]interface{}{{"id": "old"}}, Inserts: []map[string]interface{}{ {"id": "new", "document": "hello", "metadata.kind": "demo", "score": 9}, }, }) if err != nil { t.Fatalf("ApplyChanges failed: %v", err) } if ids := anySlice(deleteBody["ids"]); len(ids) != 1 || ids[0] != "old" { t.Fatalf("delete body = %#v", deleteBody) } if ids := anySlice(upsertBody["ids"]); len(ids) != 1 || ids[0] != "new" { t.Fatalf("upsert body = %#v", upsertBody) } metas := anySlice(upsertBody["metadatas"]) if len(metas) != 1 { t.Fatalf("metadatas = %#v", upsertBody["metadatas"]) } meta, _ := metas[0].(map[string]interface{}) if meta["kind"] != "demo" || meta["score"] == nil { t.Fatalf("metadata = %#v", meta) } } func TestChromaLiveSmoke(t *testing.T) { serverURL := strings.TrimSpace(os.Getenv("GONAVI_CHROMA_TEST_URL")) if serverURL == "" { t.Skip("set GONAVI_CHROMA_TEST_URL to run live Chroma smoke test") } db := newTestChromaDB(t, serverURL) collection := "gonavi_smoke_live" _, _ = db.Exec(fmt.Sprintf(`{"delete_collection":%q}`, collection)) if _, err := db.Exec(fmt.Sprintf(`{"create_collection":%q,"get_or_create":true}`, collection)); err != nil { t.Fatalf("create live collection: %v", err) } t.Cleanup(func() { _, _ = db.Exec(fmt.Sprintf(`{"delete_collection":%q}`, collection)) }) if err := db.ApplyChanges(collection, connection.ChangeSet{ Inserts: []map[string]interface{}{{ "id": "doc-1", "document": "GoNavi Chroma live smoke", "metadata.kind": "smoke", "embedding": []float64{0.1, 0.2, 0.3}, }}, }); err != nil { t.Fatalf("upsert live row: %v", err) } rows, columns, err := db.Query(fmt.Sprintf(`SELECT * FROM "%s" LIMIT 5`, collection)) if err != nil { t.Fatalf("select live rows: %v", err) } if len(rows) == 0 || rows[0]["id"] != "doc-1" || rows[0]["metadata.kind"] != "smoke" { t.Fatalf("live rows = %#v", rows) } if !containsString(columns, "metadata.kind") { t.Fatalf("live columns missing metadata.kind: %v", columns) } queryRows, queryColumns, err := db.Query(fmt.Sprintf(`{"query":%q,"query_embeddings":[[0.1,0.2,0.3]],"n_results":1}`, collection)) if err != nil { t.Fatalf("query live rows: %v", err) } if len(queryRows) == 0 || queryRows[0]["id"] != "doc-1" || !containsString(queryColumns, "distance") { t.Fatalf("live query rows = %#v columns = %v", queryRows, queryColumns) } } func containsString(items []string, target string) bool { for _, item := range items { if item == target { return true } } return false }