Merge pull request #20 from philippgille/add-unit-tests

Add unit tests

Author: philippgille

collection_test.go

@@ -0,0 +1,33 @@
+package chromem_test
+
+import (
+	"context"
+	"testing"
+
+	"github.com/philippgille/chromem-go"
+)
+
+func TestCollection_Add(t *testing.T) {
+	// Create collection
+	db := chromem.NewDB()
+	name := "test"
+	metadata := map[string]string{"foo": "bar"}
+	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
+		return []float32{-0.1, 0.1, 0.2}, nil
+	}
+	c := db.CreateCollection(name, metadata, embeddingFunc)
+	if c == nil {
+		t.Error("expected collection, got nil")
+	}
+
+	// Add document
+	ids := []string{"1", "2"}
+	metadatas := []map[string]string{{"foo": "bar"}, {"a": "b"}}
+	documents := []string{"hello world", "hallo welt"}
+	err := c.Add(context.Background(), ids, nil, metadatas, documents)
+	if err != nil {
+		t.Error("expected nil, got", err)
+	}
+
+	// TODO: Check expectations when documents become accessible
+}

db_test.go

@@ -7,6 +7,28 @@ import (
 	"github.com/philippgille/chromem-go"
 )
 
+func TestDB_CreateCollection(t *testing.T) {
+	// Values in the collection
+	name := "test"
+	metadata := map[string]string{"foo": "bar"}
+	embeddingFunc := func(_ context.Context, _ string) ([]float32, error) {
+		return []float32{-0.1, 0.1, 0.2}, nil
+	}
+
+	// Create collection
+	db := chromem.NewDB()
+	c := db.CreateCollection(name, metadata, embeddingFunc)
+	if c == nil {
+		t.Error("expected collection, got nil")
+	}
+
+	// Check expectations
+	if c.Name != name {
+		t.Error("expected name", name, "got", c.Name)
+	}
+	// TODO: Check metadata etc when they become accessible
+}
+
 func TestDB_ListCollections(t *testing.T) {
 	// Values in the collection
 	name := "test"

embed_openai_test.go

@@ -0,0 +1,86 @@
+package chromem_test
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"slices"
+	"strings"
+	"testing"
+
+	"github.com/philippgille/chromem-go"
+)
+
+type openAIResponse struct {
+	Data []struct {
+		Embedding []float32 `json:"embedding"`
+	} `json:"data"`
+}
+
+func TestNewEmbeddingFuncOpenAICompat(t *testing.T) {
+	apiKey := "secret"
+	model := "model-small"
+	baseURLSuffix := "/v1"
+	document := "hello world"
+
+	wantBody, err := json.Marshal(map[string]string{
+		"input": document,
+		"model": model,
+	})
+	if err != nil {
+		t.Error("unexpected error:", err)
+	}
+	wantRes := []float32{-0.1, 0.1, 0.2}
+
+	// Mock server
+	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		// Check URL
+		if !strings.HasSuffix(r.URL.Path, baseURLSuffix+"/embeddings") {
+			t.Error("expected URL", baseURLSuffix+"/embedding", "got", r.URL.Path)
+		}
+		// Check method
+		if r.Method != "POST" {
+			t.Error("expected method POST, got", r.Method)
+		}
+		// Check headers
+		if r.Header.Get("Authorization") != "Bearer "+apiKey {
+			t.Error("expected Authorization header", "Bearer "+apiKey, "got", r.Header.Get("Authorization"))
+		}
+		if r.Header.Get("Content-Type") != "application/json" {
+			t.Error("expected Content-Type header", "application/json", "got", r.Header.Get("Content-Type"))
+		}
+		// Check body
+		body, err := io.ReadAll(r.Body)
+		if err != nil {
+			t.Error("unexpected error:", err)
+		}
+		if !bytes.Equal(body, wantBody) {
+			t.Error("expected body", wantBody, "got", body)
+		}
+
+		// Write response
+		resp := openAIResponse{
+			Data: []struct {
+				Embedding []float32 `json:"embedding"`
+			}{
+				{Embedding: wantRes},
+			},
+		}
+		w.WriteHeader(http.StatusOK)
+		_ = json.NewEncoder(w).Encode(resp)
+	}))
+	defer ts.Close()
+	baseURL := ts.URL + baseURLSuffix
+
+	f := chromem.NewEmbeddingFuncOpenAICompat(baseURL, apiKey, model)
+	res, err := f(context.Background(), document)
+	if err != nil {
+		t.Error("expected nil, got", err)
+	}
+	if slices.Compare[[]float32](wantRes, res) != 0 {
+		t.Error("expected res", wantRes, "got", res)
+	}
+}