From fa16f626eaf8dd9c34eb023297bc55b514525954 Mon Sep 17 00:00:00 2001
From: Trayan Azarov <trayan.azarov@amikos.tech>
Date: Sat, 6 Jan 2024 12:47:33 +0200
Subject: [PATCH 1/6] feat: Auth support

- Added support for Basic and API token

Waiting on amikos-tech/chromadb-chart#39 to be implemented to add the integration tests.

Refs: #2
---
 chroma.go                  | 122 ++++++++++++++++++++++++++++++++++++-
 test/chroma_client_test.go | 107 +++++++++++++++++++++++++++++++-
 2 files changed, 226 insertions(+), 3 deletions(-)

diff --git a/chroma.go b/chroma.go
index 9ccc01a..572bce4 100644
--- a/chroma.go
+++ b/chroma.go
@@ -58,11 +58,129 @@ type Client struct {
 	ApiClient *openapiclient.APIClient //nolint
 }
 
-func NewClient(basePath string) *Client {
+type AuthType string
+
+const (
+	BASIC              AuthType = "basic"
+	TokenAuthorization AuthType = "authorization"
+	TokenXChromaToken  AuthType = "xchromatoken"
+)
+
+type AuthMethod interface {
+	GetCredentials() map[string]string
+	GetType() AuthType
+}
+
+type BasicAuth struct {
+	Username string
+	Password string
+}
+
+func (b BasicAuth) GetCredentials() map[string]string {
+	return map[string]string{
+		"username": b.Username,
+		"password": b.Password,
+	}
+}
+
+func (b BasicAuth) GetType() AuthType {
+	return BASIC
+}
+
+func NewBasicAuth(username string, password string) ClientAuthCredentials {
+	return ClientAuthCredentials{
+		AuthMethod: BasicAuth{
+			Username: username,
+			Password: password,
+		},
+	}
+}
+
+type AuthorizationTokenAuth struct {
+	Token string
+}
+
+func (t AuthorizationTokenAuth) GetType() AuthType {
+	return TokenAuthorization
+}
+
+func (t AuthorizationTokenAuth) GetCredentials() map[string]string {
+	return map[string]string{
+		"Authorization": "Bearer " + t.Token,
+	}
+}
+
+type XChromaTokenAuth struct {
+	Token string
+}
+
+func (t XChromaTokenAuth) GetType() AuthType {
+	return TokenXChromaToken
+}
+
+func (t XChromaTokenAuth) GetCredentials() map[string]string {
+	return map[string]string{
+		"X-Chroma-Token": t.Token,
+	}
+}
+
+type ClientAuthCredentials struct {
+	AuthMethod AuthMethod
+}
+
+func NewTokenAuth(token string, authType AuthType) ClientAuthCredentials {
+	switch {
+	case authType == TokenAuthorization:
+		return ClientAuthCredentials{
+			AuthMethod: AuthorizationTokenAuth{
+				Token: token,
+			},
+		}
+	case authType == TokenXChromaToken:
+		return ClientAuthCredentials{
+			AuthMethod: XChromaTokenAuth{
+				Token: token,
+			},
+		}
+	default:
+		panic("Invalid auth type")
+	}
+}
+
+type ClientConfig struct {
+	BasePath              string
+	DefaultHeaders        *map[string]string
+	ClientAuthCredentials *ClientAuthCredentials
+}
+
+func NewClientConfig(basePath string, defaultHeaders *map[string]string, clientAuthCredentials *ClientAuthCredentials) ClientConfig {
+	return ClientConfig{
+		BasePath:              basePath,
+		DefaultHeaders:        defaultHeaders,
+		ClientAuthCredentials: clientAuthCredentials,
+	}
+}
+
+func NewClient(config ClientConfig) *Client {
 	configuration := openapiclient.NewConfiguration()
+	if config.ClientAuthCredentials != nil {
+		// combine config.DefaultHeaders and config.AuthMethod.GetCredentials() maps
+		var headers = make(map[string]string)
+		if config.DefaultHeaders != nil {
+			for k, v := range *config.DefaultHeaders {
+				headers[k] = v
+			}
+		}
+		for k, v := range config.ClientAuthCredentials.AuthMethod.GetCredentials() {
+			headers[k] = v
+		}
+		configuration.DefaultHeader = headers
+	} else if config.DefaultHeaders != nil {
+		configuration.DefaultHeader = *config.DefaultHeaders
+	}
 	configuration.Servers = openapiclient.ServerConfigurations{
 		{
-			URL:         basePath,
+			URL:         config.BasePath,
 			Description: "No description provided",
 		},
 	}
diff --git a/test/chroma_client_test.go b/test/chroma_client_test.go
index b6a96a3..9e2f661 100644
--- a/test/chroma_client_test.go
+++ b/test/chroma_client_test.go
@@ -26,7 +26,9 @@ func Test_chroma_client(t *testing.T) {
 	if chromaURL == "" {
 		chromaURL = "http://localhost:8000"
 	}
-	client := chroma.NewClient(chromaURL)
+
+	clientConfig := chroma.NewClientConfig(chromaURL, nil, nil)
+	client := chroma.NewClient(clientConfig)
 
 	t.Run("Test Heartbeat", func(t *testing.T) {
 		resp, err := client.Heartbeat()
@@ -746,3 +748,106 @@ func Test_chroma_client(t *testing.T) {
 		require.Nil(t, addError)
 	})
 }
+
+func Test_chroma_client_with_basic(t *testing.T) {
+	chromaURL := os.Getenv("CHROMA_URL")
+	if chromaURL == "" {
+		chromaURL = "http://localhost:8003"
+	}
+	clientAuth := chroma.NewBasicAuth("test", "test")
+
+	clientConfig := chroma.NewClientConfig(chromaURL, nil, &clientAuth)
+	client := chroma.NewClient(clientConfig)
+
+	t.Run("Test Heartbeat", func(t *testing.T) {
+		resp, err := client.Heartbeat()
+
+		require.Nil(t, err)
+		require.NotNil(t, resp)
+		assert.Truef(t, resp["nanosecond heartbeat"] > 0, "Heartbeat should be greater than 0")
+	})
+}
+
+func Test_chroma_client_with_authorization_token(t *testing.T) {
+	chromaURL := os.Getenv("CHROMA_URL")
+	if chromaURL == "" {
+		chromaURL = "http://localhost:8001"
+	}
+	clientAuth := chroma.NewTokenAuth("test", chroma.TokenAuthorization)
+
+	clientConfig := chroma.NewClientConfig(chromaURL, nil, &clientAuth)
+	client := chroma.NewClient(clientConfig)
+
+	t.Run("Test List Collections", func(t *testing.T) {
+		collectionName1 := "test-collection1"
+		collectionName2 := "test-collection2"
+		metadata := map[string]string{}
+		apiKey := os.Getenv("OPENAI_API_KEY")
+		if apiKey == "" {
+			err := godotenv.Load("../.env")
+			if err != nil {
+				assert.Failf(t, "Error loading .env file", "%s", err)
+			}
+			apiKey = os.Getenv("OPENAI_API_KEY")
+		}
+		embeddingFunction := openai.NewOpenAIEmbeddingFunction(apiKey)
+		distanceFunction := chroma.L2
+		_, errRest := client.Reset()
+		if errRest != nil {
+			assert.Fail(t, fmt.Sprintf("Error resetting database: %s", errRest))
+		}
+		_, _ = client.CreateCollection(collectionName1, chroma.MapToAPI(metadata), true, embeddingFunction, distanceFunction)
+		_, _ = client.CreateCollection(collectionName2, chroma.MapToAPI(metadata), true, embeddingFunction, distanceFunction)
+		collections, gcerr := client.ListCollections()
+		require.Nil(t, gcerr)
+		assert.Equal(t, 2, len(collections))
+		names := make([]string, len(collections))
+		for i, person := range collections {
+			names[i] = person.Name
+		}
+		assert.Contains(t, names, collectionName1)
+		assert.Contains(t, names, collectionName2)
+	})
+}
+
+func Test_chroma_client_with_x_token(t *testing.T) {
+	chromaURL := os.Getenv("CHROMA_URL")
+	if chromaURL == "" {
+		chromaURL = "http://localhost:8002"
+	}
+	clientAuth := chroma.NewTokenAuth("test", chroma.TokenXChromaToken)
+
+	clientConfig := chroma.NewClientConfig(chromaURL, nil, &clientAuth)
+	client := chroma.NewClient(clientConfig)
+
+	t.Run("Test List Collections", func(t *testing.T) {
+		collectionName1 := "test-collection1"
+		collectionName2 := "test-collection2"
+		metadata := map[string]string{}
+		apiKey := os.Getenv("OPENAI_API_KEY")
+		if apiKey == "" {
+			err := godotenv.Load("../.env")
+			if err != nil {
+				assert.Failf(t, "Error loading .env file", "%s", err)
+			}
+			apiKey = os.Getenv("OPENAI_API_KEY")
+		}
+		embeddingFunction := openai.NewOpenAIEmbeddingFunction(apiKey)
+		distanceFunction := chroma.L2
+		_, errRest := client.Reset()
+		if errRest != nil {
+			assert.Fail(t, fmt.Sprintf("Error resetting database: %s", errRest))
+		}
+		_, _ = client.CreateCollection(collectionName1, chroma.MapToAPI(metadata), true, embeddingFunction, distanceFunction)
+		_, _ = client.CreateCollection(collectionName2, chroma.MapToAPI(metadata), true, embeddingFunction, distanceFunction)
+		collections, gcerr := client.ListCollections()
+		require.Nil(t, gcerr)
+		assert.Equal(t, 2, len(collections))
+		names := make([]string, len(collections))
+		for i, person := range collections {
+			names[i] = person.Name
+		}
+		assert.Contains(t, names, collectionName1)
+		assert.Contains(t, names, collectionName2)
+	})
+}

From e15b01417693d16378ebf20210c2cdba5742adc0 Mon Sep 17 00:00:00 2001
From: Trayan Azarov <trayan.azarov@amikos.tech>
Date: Sun, 7 Jan 2024 14:32:46 +0200
Subject: [PATCH 2/6] feat: Conditional test executions

Refs: #2
---
 test/chroma_client_test.go | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/test/chroma_client_test.go b/test/chroma_client_test.go
index 9e2f661..c70a478 100644
--- a/test/chroma_client_test.go
+++ b/test/chroma_client_test.go
@@ -22,6 +22,9 @@ import (
 )
 
 func Test_chroma_client(t *testing.T) {
+	if os.Getenv("CHROMA_TEST_AUTH_BASIC") == "TRUE" || os.Getenv("CHROMA_TEST_AUTH_AUTHORIZATION_TOKEN") == "TRUE" || os.Getenv("CHROMA_TEST_AUTH_X_TOKEN") == "TRUE" {
+		t.Skip("Skipping Chroma Client tests during auth tests.")
+	}
 	chromaURL := os.Getenv("CHROMA_URL")
 	if chromaURL == "" {
 		chromaURL = "http://localhost:8000"
@@ -750,6 +753,9 @@ func Test_chroma_client(t *testing.T) {
 }
 
 func Test_chroma_client_with_basic(t *testing.T) {
+	if os.Getenv("CHROMA_TEST_AUTH_BASIC") != "TRUE" {
+		t.Skip("Skipping Test_chroma_client_with_basic")
+	}
 	chromaURL := os.Getenv("CHROMA_URL")
 	if chromaURL == "" {
 		chromaURL = "http://localhost:8003"
@@ -769,6 +775,9 @@ func Test_chroma_client_with_basic(t *testing.T) {
 }
 
 func Test_chroma_client_with_authorization_token(t *testing.T) {
+	if os.Getenv("CHROMA_TEST_AUTH_AUTHORIZATION_TOKEN") != "TRUE" {
+		t.Skip("Skipping Test_chroma_client_with_authorization_token")
+	}
 	chromaURL := os.Getenv("CHROMA_URL")
 	if chromaURL == "" {
 		chromaURL = "http://localhost:8001"
@@ -811,6 +820,9 @@ func Test_chroma_client_with_authorization_token(t *testing.T) {
 }
 
 func Test_chroma_client_with_x_token(t *testing.T) {
+	if os.Getenv("CHROMA_TEST_AUTH_X_TOKEN") != "TRUE" {
+		t.Skip("Skipping Test_chroma_client_with_x_token")
+	}
 	chromaURL := os.Getenv("CHROMA_URL")
 	if chromaURL == "" {
 		chromaURL = "http://localhost:8002"

From 5695e3a5dfe730b973aecc369d29ff2ae6d9f5a7 Mon Sep 17 00:00:00 2001
From: Trayan Azarov <trayan.azarov@amikos.tech>
Date: Sat, 6 Jan 2024 12:47:33 +0200
Subject: [PATCH 3/6] feat: Auth support

- Added support for Basic and API token

Waiting on amikos-tech/chromadb-chart#39 to be implemented to add the integration tests.

Refs: #2
---
 chroma.go                  | 122 ++++++++++++++++++++++++++++++++++++-
 test/chroma_client_test.go | 107 +++++++++++++++++++++++++++++++-
 2 files changed, 226 insertions(+), 3 deletions(-)

diff --git a/chroma.go b/chroma.go
index 9ccc01a..572bce4 100644
--- a/chroma.go
+++ b/chroma.go
@@ -58,11 +58,129 @@ type Client struct {
 	ApiClient *openapiclient.APIClient //nolint
 }
 
-func NewClient(basePath string) *Client {
+type AuthType string
+
+const (
+	BASIC              AuthType = "basic"
+	TokenAuthorization AuthType = "authorization"
+	TokenXChromaToken  AuthType = "xchromatoken"
+)
+
+type AuthMethod interface {
+	GetCredentials() map[string]string
+	GetType() AuthType
+}
+
+type BasicAuth struct {
+	Username string
+	Password string
+}
+
+func (b BasicAuth) GetCredentials() map[string]string {
+	return map[string]string{
+		"username": b.Username,
+		"password": b.Password,
+	}
+}
+
+func (b BasicAuth) GetType() AuthType {
+	return BASIC
+}
+
+func NewBasicAuth(username string, password string) ClientAuthCredentials {
+	return ClientAuthCredentials{
+		AuthMethod: BasicAuth{
+			Username: username,
+			Password: password,
+		},
+	}
+}
+
+type AuthorizationTokenAuth struct {
+	Token string
+}
+
+func (t AuthorizationTokenAuth) GetType() AuthType {
+	return TokenAuthorization
+}
+
+func (t AuthorizationTokenAuth) GetCredentials() map[string]string {
+	return map[string]string{
+		"Authorization": "Bearer " + t.Token,
+	}
+}
+
+type XChromaTokenAuth struct {
+	Token string
+}
+
+func (t XChromaTokenAuth) GetType() AuthType {
+	return TokenXChromaToken
+}
+
+func (t XChromaTokenAuth) GetCredentials() map[string]string {
+	return map[string]string{
+		"X-Chroma-Token": t.Token,
+	}
+}
+
+type ClientAuthCredentials struct {
+	AuthMethod AuthMethod
+}
+
+func NewTokenAuth(token string, authType AuthType) ClientAuthCredentials {
+	switch {
+	case authType == TokenAuthorization:
+		return ClientAuthCredentials{
+			AuthMethod: AuthorizationTokenAuth{
+				Token: token,
+			},
+		}
+	case authType == TokenXChromaToken:
+		return ClientAuthCredentials{
+			AuthMethod: XChromaTokenAuth{
+				Token: token,
+			},
+		}
+	default:
+		panic("Invalid auth type")
+	}
+}
+
+type ClientConfig struct {
+	BasePath              string
+	DefaultHeaders        *map[string]string
+	ClientAuthCredentials *ClientAuthCredentials
+}
+
+func NewClientConfig(basePath string, defaultHeaders *map[string]string, clientAuthCredentials *ClientAuthCredentials) ClientConfig {
+	return ClientConfig{
+		BasePath:              basePath,
+		DefaultHeaders:        defaultHeaders,
+		ClientAuthCredentials: clientAuthCredentials,
+	}
+}
+
+func NewClient(config ClientConfig) *Client {
 	configuration := openapiclient.NewConfiguration()
+	if config.ClientAuthCredentials != nil {
+		// combine config.DefaultHeaders and config.AuthMethod.GetCredentials() maps
+		var headers = make(map[string]string)
+		if config.DefaultHeaders != nil {
+			for k, v := range *config.DefaultHeaders {
+				headers[k] = v
+			}
+		}
+		for k, v := range config.ClientAuthCredentials.AuthMethod.GetCredentials() {
+			headers[k] = v
+		}
+		configuration.DefaultHeader = headers
+	} else if config.DefaultHeaders != nil {
+		configuration.DefaultHeader = *config.DefaultHeaders
+	}
 	configuration.Servers = openapiclient.ServerConfigurations{
 		{
-			URL:         basePath,
+			URL:         config.BasePath,
 			Description: "No description provided",
 		},
 	}
diff --git a/test/chroma_client_test.go b/test/chroma_client_test.go
index b6a96a3..9e2f661 100644
--- a/test/chroma_client_test.go
+++ b/test/chroma_client_test.go
@@ -26,7 +26,9 @@ func Test_chroma_client(t *testing.T) {
 	if chromaURL == "" {
 		chromaURL = "http://localhost:8000"
 	}
-	client := chroma.NewClient(chromaURL)
+
+	clientConfig := chroma.NewClientConfig(chromaURL, nil, nil)
+	client := chroma.NewClient(clientConfig)
 
 	t.Run("Test Heartbeat", func(t *testing.T) {
 		resp, err := client.Heartbeat()
@@ -746,3 +748,106 @@ func Test_chroma_client(t *testing.T) {
 		require.Nil(t, addError)
 	})
 }
+
+func Test_chroma_client_with_basic(t *testing.T) {
+	chromaURL := os.Getenv("CHROMA_URL")
+	if chromaURL == "" {
+		chromaURL = "http://localhost:8003"
+	}
+	clientAuth := chroma.NewBasicAuth("test", "test")
+
+	clientConfig := chroma.NewClientConfig(chromaURL, nil, &clientAuth)
+	client := chroma.NewClient(clientConfig)
+
+	t.Run("Test Heartbeat", func(t *testing.T) {
+		resp, err := client.Heartbeat()
+
+		require.Nil(t, err)
+		require.NotNil(t, resp)
+		assert.Truef(t, resp["nanosecond heartbeat"] > 0, "Heartbeat should be greater than 0")
+	})
+}
+
+func Test_chroma_client_with_authorization_token(t *testing.T) {
+	chromaURL := os.Getenv("CHROMA_URL")
+	if chromaURL == "" {
+		chromaURL = "http://localhost:8001"
+	}
+	clientAuth := chroma.NewTokenAuth("test", chroma.TokenAuthorization)
+
+	clientConfig := chroma.NewClientConfig(chromaURL, nil, &clientAuth)
+	client := chroma.NewClient(clientConfig)
+
+	t.Run("Test List Collections", func(t *testing.T) {
+		collectionName1 := "test-collection1"
+		collectionName2 := "test-collection2"
+		metadata := map[string]string{}
+		apiKey := os.Getenv("OPENAI_API_KEY")
+		if apiKey == "" {
+			err := godotenv.Load("../.env")
+			if err != nil {
+				assert.Failf(t, "Error loading .env file", "%s", err)
+			}
+			apiKey = os.Getenv("OPENAI_API_KEY")
+		}
+		embeddingFunction := openai.NewOpenAIEmbeddingFunction(apiKey)
+		distanceFunction := chroma.L2
+		_, errRest := client.Reset()
+		if errRest != nil {
+			assert.Fail(t, fmt.Sprintf("Error resetting database: %s", errRest))
+		}
+		_, _ = client.CreateCollection(collectionName1, chroma.MapToAPI(metadata), true, embeddingFunction, distanceFunction)
+		_, _ = client.CreateCollection(collectionName2, chroma.MapToAPI(metadata), true, embeddingFunction, distanceFunction)
+		collections, gcerr := client.ListCollections()
+		require.Nil(t, gcerr)
+		assert.Equal(t, 2, len(collections))
+		names := make([]string, len(collections))
+		for i, person := range collections {
+			names[i] = person.Name
+		}
+		assert.Contains(t, names, collectionName1)
+		assert.Contains(t, names, collectionName2)
+	})
+}
+
+func Test_chroma_client_with_x_token(t *testing.T) {
+	chromaURL := os.Getenv("CHROMA_URL")
+	if chromaURL == "" {
+		chromaURL = "http://localhost:8002"
+	}
+	clientAuth := chroma.NewTokenAuth("test", chroma.TokenXChromaToken)
+
+	clientConfig := chroma.NewClientConfig(chromaURL, nil, &clientAuth)
+	client := chroma.NewClient(clientConfig)
+
+	t.Run("Test List Collections", func(t *testing.T) {
+		collectionName1 := "test-collection1"
+		collectionName2 := "test-collection2"
+		metadata := map[string]string{}
+		apiKey := os.Getenv("OPENAI_API_KEY")
+		if apiKey == "" {
+			err := godotenv.Load("../.env")
+			if err != nil {
+				assert.Failf(t, "Error loading .env file", "%s", err)
+			}
+			apiKey = os.Getenv("OPENAI_API_KEY")
+		}
+		embeddingFunction := openai.NewOpenAIEmbeddingFunction(apiKey)
+		distanceFunction := chroma.L2
+		_, errRest := client.Reset()
+		if errRest != nil {
+			assert.Fail(t, fmt.Sprintf("Error resetting database: %s", errRest))
+		}
+		_, _ = client.CreateCollection(collectionName1, chroma.MapToAPI(metadata), true, embeddingFunction, distanceFunction)
+		_, _ = client.CreateCollection(collectionName2, chroma.MapToAPI(metadata), true, embeddingFunction, distanceFunction)
+		collections, gcerr := client.ListCollections()
+		require.Nil(t, gcerr)
+		assert.Equal(t, 2, len(collections))
+		names := make([]string, len(collections))
+		for i, person := range collections {
+			names[i] = person.Name
+		}
+		assert.Contains(t, names, collectionName1)
+		assert.Contains(t, names, collectionName2)
+	})
+}

From 684d430fac71880b340ac0293f46c602d132d8a4 Mon Sep 17 00:00:00 2001
From: Trayan Azarov <trayan.azarov@amikos.tech>
Date: Sun, 7 Jan 2024 14:32:46 +0200
Subject: [PATCH 4/6] feat: Conditional test executions

Refs: #2
---
 test/chroma_client_test.go | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/test/chroma_client_test.go b/test/chroma_client_test.go
index 9e2f661..c70a478 100644
--- a/test/chroma_client_test.go
+++ b/test/chroma_client_test.go
@@ -22,6 +22,9 @@ import (
 )
 
 func Test_chroma_client(t *testing.T) {
+	if os.Getenv("CHROMA_TEST_AUTH_BASIC") == "TRUE" || os.Getenv("CHROMA_TEST_AUTH_AUTHORIZATION_TOKEN") == "TRUE" || os.Getenv("CHROMA_TEST_AUTH_X_TOKEN") == "TRUE" {
+		t.Skip("Skipping Chroma Client tests during auth tests.")
+	}
 	chromaURL := os.Getenv("CHROMA_URL")
 	if chromaURL == "" {
 		chromaURL = "http://localhost:8000"
@@ -750,6 +753,9 @@ func Test_chroma_client(t *testing.T) {
 }
 
 func Test_chroma_client_with_basic(t *testing.T) {
+	if os.Getenv("CHROMA_TEST_AUTH_BASIC") != "TRUE" {
+		t.Skip("Skipping Test_chroma_client_with_basic")
+	}
 	chromaURL := os.Getenv("CHROMA_URL")
 	if chromaURL == "" {
 		chromaURL = "http://localhost:8003"
@@ -769,6 +775,9 @@ func Test_chroma_client_with_basic(t *testing.T) {
 }
 
 func Test_chroma_client_with_authorization_token(t *testing.T) {
+	if os.Getenv("CHROMA_TEST_AUTH_AUTHORIZATION_TOKEN") != "TRUE" {
+		t.Skip("Skipping Test_chroma_client_with_authorization_token")
+	}
 	chromaURL := os.Getenv("CHROMA_URL")
 	if chromaURL == "" {
 		chromaURL = "http://localhost:8001"
@@ -811,6 +820,9 @@ func Test_chroma_client_with_authorization_token(t *testing.T) {
 }
 
 func Test_chroma_client_with_x_token(t *testing.T) {
+	if os.Getenv("CHROMA_TEST_AUTH_X_TOKEN") != "TRUE" {
+		t.Skip("Skipping Test_chroma_client_with_x_token")
+	}
 	chromaURL := os.Getenv("CHROMA_URL")
 	if chromaURL == "" {
 		chromaURL = "http://localhost:8002"

From b5cdd2c2606a1342dd2e91f83f5cb3daf54a329b Mon Sep 17 00:00:00 2001
From: Trayan Azarov <trayan.azarov@amikos.tech>
Date: Tue, 9 Jan 2024 15:03:55 +0200
Subject: [PATCH 5/6] fix: Fixed for context

Refs: #2
---
 test/chroma_client_test.go | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/test/chroma_client_test.go b/test/chroma_client_test.go
index 325623c..4b9300b 100644
--- a/test/chroma_client_test.go
+++ b/test/chroma_client_test.go
@@ -767,7 +767,7 @@ func Test_chroma_client_with_basic(t *testing.T) {
 	client := chroma.NewClient(clientConfig)
 
 	t.Run("Test Heartbeat", func(t *testing.T) {
-		resp, err := client.Heartbeat()
+		resp, err := client.Heartbeat(context.Background())
 
 		require.Nil(t, err)
 		require.NotNil(t, resp)
@@ -802,13 +802,13 @@ func Test_chroma_client_with_authorization_token(t *testing.T) {
 		}
 		embeddingFunction := openai.NewOpenAIEmbeddingFunction(apiKey)
 		distanceFunction := chroma.L2
-		_, errRest := client.Reset()
+		_, errRest := client.Reset(context.Background())
 		if errRest != nil {
 			assert.Fail(t, fmt.Sprintf("Error resetting database: %s", errRest))
 		}
-		_, _ = client.CreateCollection(collectionName1, chroma.MapToAPI(metadata), true, embeddingFunction, distanceFunction)
-		_, _ = client.CreateCollection(collectionName2, chroma.MapToAPI(metadata), true, embeddingFunction, distanceFunction)
-		collections, gcerr := client.ListCollections()
+		_, _ = client.CreateCollection(context.Background(), collectionName1, chroma.MapToAPI(metadata), true, embeddingFunction, distanceFunction)
+		_, _ = client.CreateCollection(context.Background(), collectionName2, chroma.MapToAPI(metadata), true, embeddingFunction, distanceFunction)
+		collections, gcerr := client.ListCollections(context.Background())
 		require.Nil(t, gcerr)
 		assert.Equal(t, 2, len(collections))
 		names := make([]string, len(collections))
@@ -847,13 +847,13 @@ func Test_chroma_client_with_x_token(t *testing.T) {
 		}
 		embeddingFunction := openai.NewOpenAIEmbeddingFunction(apiKey)
 		distanceFunction := chroma.L2
-		_, errRest := client.Reset()
+		_, errRest := client.Reset(context.Background())
 		if errRest != nil {
 			assert.Fail(t, fmt.Sprintf("Error resetting database: %s", errRest))
 		}
-		_, _ = client.CreateCollection(collectionName1, chroma.MapToAPI(metadata), true, embeddingFunction, distanceFunction)
-		_, _ = client.CreateCollection(collectionName2, chroma.MapToAPI(metadata), true, embeddingFunction, distanceFunction)
-		collections, gcerr := client.ListCollections()
+		_, _ = client.CreateCollection(context.Background(), collectionName1, chroma.MapToAPI(metadata), true, embeddingFunction, distanceFunction)
+		_, _ = client.CreateCollection(context.Background(), collectionName2, chroma.MapToAPI(metadata), true, embeddingFunction, distanceFunction)
+		collections, gcerr := client.ListCollections(context.Background())
 		require.Nil(t, gcerr)
 		assert.Equal(t, 2, len(collections))
 		names := make([]string, len(collections))

From ac749ee2d7b4e33a402f214076903f7241e9aa55 Mon Sep 17 00:00:00 2001
From: Trayan Azarov <trayan.azarov@amikos.tech>
Date: Tue, 9 Jan 2024 22:36:14 +0200
Subject: [PATCH 6/6] feat: Added CI tests

Refs: #2
---
 .github/workflows/go.yml   | 163 +++++++++++++++++++++++++------------
 test/chroma_client_test.go |   8 +-
 2 files changed, 115 insertions(+), 56 deletions(-)

diff --git a/.github/workflows/go.yml b/.github/workflows/go.yml
index 18545c8..b52d38e 100644
--- a/.github/workflows/go.yml
+++ b/.github/workflows/go.yml
@@ -9,62 +9,121 @@ on:
   pull_request:
     branches: [ "main" ]
 
+env:
+  BASIC_AUTH_NS: chroma-auth-basic
+  TOKEN_BEARER_AUTH_NS: chroma-auth-token-bearer
+  TOKEN_XTOKEN_AUTH_NS: chroma-auth-token-xtoken
 jobs:
-
   build:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-        chroma-version: [0.4.8, 0.4.20]
+        chroma-version: [ 0.4.8, 0.4.20 ]
     environment: Test
     steps:
-    - uses: actions/checkout@v3
-    - name: Set up Go
-      uses: actions/setup-go@v4
-      with:
-        go-version-file: 'go.mod'
-    - name: Run golangci-lint
-      uses: golangci/golangci-lint-action@v3
-      with:
-        version: latest
-    - name: Build
-      run: make build
-    - name: Install Helm
-      uses: azure/setup-helm@v1
-      with:
-        version: v3.4.0
-    - name: start minikube
-      id: minikube
-      uses: medyagh/setup-minikube@latest
-      with:
-        kubernetes-version: 1.27.3
-    - name: Add helm repo
-      run: |
-        set -e
-        helm repo add chromadb https://amikos-tech.github.io/chromadb-chart/
-        helm repo update
-    - name: Install chromadb
-      run: |
-        set -e
-        helm install chromadb chromadb/chromadb --set chromadb.allowReset=true,chromadb.apiVersion=${{ matrix.chroma-version }},chromadb.auth.enabled=false
-    - name: Wait for deployment to be ready
-      id: wait-and-set
-      run: |
-        set -e
-        kubectl wait \
-        --for=condition=ready pod \
-        --selector=app.kubernetes.io/name=chromadb \
-        --timeout=120s
-        echo "chroma-url=$(minikube service chromadb --url)" >> $GITHUB_OUTPUT
-    - name: Hearthbeat
-      run: |
-        set -e
-        kubectl get svc -A
-        curl $(minikube service chromadb --url)/api/v1
-    - name: Test
-      run: make gotest
-      env:
-        OPENAI_API_KEY:  ${{ secrets.OPENAI_API_KEY }}
-        COHERE_API_KEY:  ${{ secrets.COHERE_API_KEY }}
-        HF_API_KEY:  ${{ secrets.HF_API_KEY }}
-        CHROMA_URL: ${{steps.wait-and-set.outputs.chroma-url}}
+      - uses: actions/checkout@v3
+      - name: Set up Go
+        uses: actions/setup-go@v4
+        with:
+          go-version-file: 'go.mod'
+      - name: Run golangci-lint
+        uses: golangci/golangci-lint-action@v3
+        with:
+          version: latest
+      - name: Build
+        run: make build
+      - name: Install Helm
+        uses: azure/setup-helm@v1
+        with:
+          version: v3.4.0
+      - name: start minikube
+        id: minikube
+        uses: medyagh/setup-minikube@latest
+        with:
+          kubernetes-version: 1.27.3
+      - name: Add helm repo
+        run: |
+          set -e
+          helm repo add chromadb https://amikos-tech.github.io/chromadb-chart/
+          helm repo update
+      - name: Install chromadb
+        run: |
+          set -e
+          helm install chromadb chromadb/chromadb --set chromadb.allowReset=true,chromadb.apiVersion=${{ matrix.chroma-version }},chromadb.auth.enabled=false
+          kubectl create namespace ${{ env.BASIC_AUTH_NS }}
+          kubectl create namespace ${{ env.TOKEN_BEARER_AUTH_NS }}
+          kubectl create namespace ${{ env.TOKEN_XTOKEN_AUTH_NS }}
+          helm install chromadb chromadb/chromadb --namespace ${{ env.BASIC_AUTH_NS }} --set chromadb.allowReset=true,chromadb.apiVersion=${{ matrix.chroma-version }},chromadb.auth.enabled=true,chromadb.auth.type=basic
+          helm install chromadb chromadb/chromadb --namespace ${{ env.TOKEN_BEARER_AUTH_NS }} --set chromadb.allowReset=true,chromadb.apiVersion=${{ matrix.chroma-version }},chromadb.auth.enabled=true,chromadb.auth.type=token,chromadb.auth.token.tokenHeader=AUTHORIZATION
+          helm install chromadb chromadb/chromadb --namespace ${{ env.TOKEN_XTOKEN_AUTH_NS }} --set chromadb.allowReset=true,chromadb.apiVersion=${{ matrix.chroma-version }},chromadb.auth.enabled=true,chromadb.auth.type=token,chromadb.auth.token.tokenHeader=X-CHROMA-TOKEN
+      - name: Wait for deployment to be ready
+        id: wait-and-set
+        run: |
+          set -e
+          kubectl wait \
+          --for=condition=ready pod \
+          --selector=app.kubernetes.io/name=chromadb \
+          --timeout=120s
+          kubectl wait \
+          --for=condition=ready pod \
+          --selector=app.kubernetes.io/name=chromadb \
+          --namespace ${{ env.BASIC_AUTH_NS }} \
+          --timeout=120s
+          kubectl wait \
+          --for=condition=ready pod \
+          --selector=app.kubernetes.io/name=chromadb \
+          --namespace ${{ env.TOKEN_BEARER_AUTH_NS }} \
+          --timeout=120s
+          kubectl wait \
+          --for=condition=ready pod \
+          --selector=app.kubernetes.io/name=chromadb \
+          --namespace ${{ env.TOKEN_XTOKEN_AUTH_NS }} \
+          --timeout=120s
+          echo "chroma-url=$(minikube service chromadb --url)" >> $GITHUB_OUTPUT
+          echo "chroma-basic-auth-url=$(minikube service chromadb --url -n ${{ env.BASIC_AUTH_NS }})" >> $GITHUB_OUTPUT
+          echo "chroma-token-auth-url=$(minikube service chromadb --url -n ${{ env.TOKEN_BEARER_AUTH_NS }})" >> $GITHUB_OUTPUT
+          echo "chroma-token-xtoken-url=$(minikube service chromadb --url -n ${{ env.TOKEN_XTOKEN_AUTH_NS }})" >> $GITHUB_OUTPUT
+          echo "chroma-basic-user=$(kubectl --namespace ${{ env.BASIC_AUTH_NS }} get secret chromadb-auth -o jsonpath="{.data.username}" | base64 --decode)" >> $GITHUB_OUTPUT
+          echo "chroma-basic-password=$(kubectl --namespace ${{ env.BASIC_AUTH_NS }} get secret chromadb-auth -o jsonpath="{.data.password}" | base64 --decode)" >> $GITHUB_OUTPUT
+          echo "chroma-bearer-token=$(kubectl --namespace ${{ env.TOKEN_BEARER_AUTH_NS }} get secret chromadb-auth -o jsonpath="{.data.token}" | base64 --decode)" >> $GITHUB_OUTPUT
+          echo "chroma-x-token=$(kubectl --namespace ${{ env.TOKEN_XTOKEN_AUTH_NS }} get secret chromadb-auth -o jsonpath="{.data.token}" | base64 --decode)" >> $GITHUB_OUTPUT
+      - name: Hearthbeat
+        run: |
+          set -e
+          kubectl get svc -A
+          curl $(minikube service chromadb --url)/api/v1
+      - name: Test
+        run: make gotest
+        env:
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
+          HF_API_KEY: ${{ secrets.HF_API_KEY }}
+          CHROMA_URL: ${{steps.wait-and-set.outputs.chroma-url}}
+      - name: Test Auth - Basic
+        run: make gotest
+        env:
+          CHROMA_AUTH_BASIC_USERNAME: ${{steps.wait-and-set.outputs.chroma-basic-user}}
+          CHROMA_AUTH_BASIC_PASSWORD: ${{steps.wait-and-set.outputs.chroma-basic-password}}
+          CHROMA_TEST_AUTH_BASIC: TRUE
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
+          HF_API_KEY: ${{ secrets.HF_API_KEY }}
+          CHROMA_URL: ${{steps.wait-and-set.outputs.chroma-basic-auth-url}}
+      - name: Test Auth - Bearer Authorization
+        run: make gotest
+        env:
+          CHROMA_AUTH_TOKEN: ${{ steps.wait-and-set.outputs.chroma-bearer-token }}
+          CHROMA_TEST_AUTH_AUTHORIZATION_TOKEN: TRUE
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
+          HF_API_KEY: ${{ secrets.HF_API_KEY }}
+          CHROMA_URL: ${{ steps.wait-and-set.outputs.chroma-token-auth-url }}
+      - name: Test Auth - Bearer Authorization
+        run: make gotest
+        env:
+          CHROMA_AUTH_TOKEN: ${{ steps.wait-and-set.outputs.chroma-x-token }}
+          CHROMA_TEST_AUTH_X_TOKEN: TRUE
+          OPENAI_API_KEY: ${{ secrets.OPENAI_API_KEY }}
+          COHERE_API_KEY: ${{ secrets.COHERE_API_KEY }}
+          HF_API_KEY: ${{ secrets.HF_API_KEY }}
+          CHROMA_URL: ${{ steps.wait-and-set.outputs.chroma-token-xtoken-url }}
diff --git a/test/chroma_client_test.go b/test/chroma_client_test.go
index 4b9300b..f02dacb 100644
--- a/test/chroma_client_test.go
+++ b/test/chroma_client_test.go
@@ -759,9 +759,9 @@ func Test_chroma_client_with_basic(t *testing.T) {
 	}
 	chromaURL := os.Getenv("CHROMA_URL")
 	if chromaURL == "" {
-		chromaURL = "http://localhost:8003"
+		chromaURL = "http://localhost:8000"
 	}
-	clientAuth := chroma.NewBasicAuth("test", "test")
+	clientAuth := chroma.NewBasicAuth(os.Getenv("CHROMA_AUTH_BASIC_USERNAME"), os.Getenv("CHROMA_AUTH_BASIC_PASSWORD"))
 
 	clientConfig := chroma.NewClientConfig(chromaURL, nil, &clientAuth)
 	client := chroma.NewClient(clientConfig)
@@ -781,9 +781,9 @@ func Test_chroma_client_with_authorization_token(t *testing.T) {
 	}
 	chromaURL := os.Getenv("CHROMA_URL")
 	if chromaURL == "" {
-		chromaURL = "http://localhost:8001"
+		chromaURL = "http://localhost:8000"
 	}
-	clientAuth := chroma.NewTokenAuth("test", chroma.TokenAuthorization)
+	clientAuth := chroma.NewTokenAuth(os.Getenv("CHROMA_AUTH_TOKEN"), chroma.TokenAuthorization)
 
 	clientConfig := chroma.NewClientConfig(chromaURL, nil, &clientAuth)
 	client := chroma.NewClient(clientConfig)