From 1a072c7b01607127a264b6dcb5159bca4843abeb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philipp=20Gill=C3=A9?= <philipp.gille@gmail.com>
Date: Sat, 2 Mar 2024 15:21:33 +0100
Subject: [PATCH 1/3] Add more validation to DB.CreateCollection

---
 db.go | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/db.go b/db.go
index 95cf337..5cdf0d9 100644
--- a/db.go
+++ b/db.go
@@ -138,6 +138,9 @@ func NewPersistentDB(path string) (*DB, error) {
 //   - embeddingFunc: Optional function to use to embed documents.
 //     Uses the default embedding function if not provided.
 func (db *DB) CreateCollection(name string, metadata map[string]string, embeddingFunc EmbeddingFunc) (*Collection, error) {
+	if name == "" {
+		return nil, errors.New("collection name is empty")
+	}
 	if embeddingFunc == nil {
 		embeddingFunc = NewEmbeddingFuncDefault()
 	}

From d7990c795eecd34e87a2493911ccdb65907798d2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philipp=20Gill=C3=A9?= <philipp.gille@gmail.com>
Date: Sat, 2 Mar 2024 15:21:37 +0100
Subject: [PATCH 2/3] Add unit test

---
 db_test.go | 35 ++++++++++++++++++++++-------------
 1 file changed, 22 insertions(+), 13 deletions(-)

diff --git a/db_test.go b/db_test.go
index 449131a..da74bad 100644
--- a/db_test.go
+++ b/db_test.go
@@ -15,21 +15,30 @@ func TestDB_CreateCollection(t *testing.T) {
 		return []float32{-0.1, 0.1, 0.2}, nil
 	}
 
-	// Create collection
 	db := chromem.NewDB()
-	c, err := db.CreateCollection(name, metadata, embeddingFunc)
-	if err != nil {
-		t.Error("expected no error, got", err)
-	}
-	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
+	t.Run("OK", func(t *testing.T) {
+		c, err := db.CreateCollection(name, metadata, embeddingFunc)
+		if err != nil {
+			t.Error("expected no error, got", err)
+		}
+		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
+	})
+
+	t.Run("NOK - Empty name", func(t *testing.T) {
+		_, err := db.CreateCollection("", metadata, embeddingFunc)
+		if err == nil {
+			t.Error("expected error, got nil")
+		}
+	})
 }
 
 func TestDB_ListCollections(t *testing.T) {

From 20c2f9a899bdeb889646350345af0feeddaec7de Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Philipp=20Gill=C3=A9?= <philipp.gille@gmail.com>
Date: Sat, 2 Mar 2024 15:21:46 +0100
Subject: [PATCH 3/3] Add more validation to Collection.Query

---
 collection.go | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/collection.go b/collection.go
index bffa921..4e45f2a 100644
--- a/collection.go
+++ b/collection.go
@@ -107,16 +107,19 @@ func (c *Collection) Count() int {
 //   - where: Conditional filtering on metadata. Optional.
 //   - whereDocument: Conditional filtering on documents. Optional.
 func (c *Collection) Query(ctx context.Context, queryText string, nResults int, where, whereDocument map[string]string) ([]Result, error) {
+	if queryText == "" {
+		return nil, errors.New("queryText is empty")
+	}
+	if nResults <= 0 {
+		return nil, errors.New("nResults must be > 0")
+	}
+
 	c.documentsLock.RLock()
 	defer c.documentsLock.RUnlock()
 	if len(c.documents) == 0 {
 		return nil, nil
 	}
 
-	if nResults <= 0 {
-		return nil, errors.New("nResults must be > 0")
-	}
-
 	// Validate whereDocument operators
 	for k := range whereDocument {
 		if !slices.Contains(supportedFilters, k) {