chore: use a map

testcontainers · Dec 13, 2024 · ccd1974 · ccd1974
1 parent 857a378
commit ccd1974
Show file tree

Hide file tree

Showing 4 changed files with 15 additions and 43 deletions.
diff --git a/modules/ollama/examples_test.go b/modules/ollama/examples_test.go
@@ -178,7 +178,7 @@ func ExampleRun_withLocal() {
 	ctx := context.Background()
 
 	// localOllama {
-	ollamaContainer, err := tcollama.Run(ctx, "ollama/ollama:0.3.13", tcollama.WithUseLocal("OLLAMA_DEBUG=true"))
+	ollamaContainer, err := tcollama.Run(ctx, "ollama/ollama:0.3.13", tcollama.WithUseLocal(map[string]string{"OLLAMA_DEBUG": "true"}))
 	defer func() {
 		if err := testcontainers.TerminateContainer(ollamaContainer); err != nil {
 			log.Printf("failed to terminate container: %s", err)

diff --git a/modules/ollama/local_test.go b/modules/ollama/local_test.go
@@ -29,7 +29,7 @@ func TestRun_local(t *testing.T) {
 	ollamaContainer, err := ollama.Run(
 		ctx,
 		"ollama/ollama:0.1.25",
-		ollama.WithUseLocal("FOO=BAR"),
+		ollama.WithUseLocal(map[string]string{"FOO": "BAR"}),
 	)
 	testcontainers.CleanupContainer(t, ollamaContainer)
 	require.NoError(t, err)
@@ -266,7 +266,7 @@ func TestRun_localWithCustomLogFile(t *testing.T) {
 
 	ctx := context.Background()
 
-	ollamaContainer, err := ollama.Run(ctx, "ollama/ollama:0.1.25", ollama.WithUseLocal("FOO=BAR"))
+	ollamaContainer, err := ollama.Run(ctx, "ollama/ollama:0.1.25", ollama.WithUseLocal(map[string]string{"FOO": "BAR"}))
 	require.NoError(t, err)
 	testcontainers.CleanupContainer(t, ollamaContainer)
 
@@ -285,7 +285,7 @@ func TestRun_localWithCustomHost(t *testing.T) {
 
 	ctx := context.Background()
 
-	ollamaContainer, err := ollama.Run(ctx, "ollama/ollama:0.1.25", ollama.WithUseLocal("FOO=BAR"))
+	ollamaContainer, err := ollama.Run(ctx, "ollama/ollama:0.1.25", ollama.WithUseLocal(nil))
 	require.NoError(t, err)
 	testcontainers.CleanupContainer(t, ollamaContainer)
 

diff --git a/modules/ollama/options.go b/modules/ollama/options.go
@@ -2,8 +2,6 @@ package ollama
 
 import (
 	"context"
-	"fmt"
-	"strings"
 
 	"github.com/docker/docker/api/types/container"
 
@@ -44,29 +42,23 @@ var _ testcontainers.ContainerCustomizer = (*UseLocal)(nil)
 
 // UseLocal will use the local Ollama instance instead of pulling the Docker image.
 type UseLocal struct {
-	env []string
+	env map[string]string
 }
 
 // WithUseLocal the module will use the local Ollama instance instead of pulling the Docker image.
 // Pass the environment variables you need to set for the Ollama binary to be used,
 // in the format of "KEY=VALUE". KeyValue pairs with the wrong format will cause an error.
-func WithUseLocal(keyVal ...string) UseLocal {
+func WithUseLocal(keyVal map[string]string) UseLocal {
 	return UseLocal{env: keyVal}
 }
 
 // Customize implements the ContainerCustomizer interface, taking the key value pairs
 // and setting them as environment variables for the Ollama binary.
 // In the case of an invalid key value pair, an error is returned.
 func (u UseLocal) Customize(req *testcontainers.GenericContainerRequest) error {
-	env := make(map[string]string)
-	for _, kv := range u.env {
-		parts := strings.SplitN(kv, "=", 2)
-		if len(parts) != 2 {
-			return fmt.Errorf("invalid environment variable: %s", kv)
-		}
-
-		env[parts[0]] = parts[1]
+	if len(u.env) == 0 {
+		return nil
 	}
 
-	return testcontainers.WithEnv(env)(req)
+	return testcontainers.WithEnv(u.env)(req)
 }
diff --git a/modules/ollama/options_test.go b/modules/ollama/options_test.go
@@ -12,38 +12,18 @@ import (
 func TestWithUseLocal(t *testing.T) {
 	req := testcontainers.GenericContainerRequest{}
 
-	t.Run("keyVal/valid", func(t *testing.T) {
-		opt := ollama.WithUseLocal("OLLAMA_MODELS=/path/to/models")
+	t.Run("empty", func(t *testing.T) {
+		opt := ollama.WithUseLocal(nil)
 		err := opt.Customize(&req)
 		require.NoError(t, err)
-		require.Equal(t, "/path/to/models", req.Env["OLLAMA_MODELS"])
-	})
-
-	t.Run("keyVal/invalid", func(t *testing.T) {
-		opt := ollama.WithUseLocal("OLLAMA_MODELS")
-		err := opt.Customize(&req)
-		require.Error(t, err)
-	})
-
-	t.Run("keyVal/valid/multiple", func(t *testing.T) {
-		opt := ollama.WithUseLocal("OLLAMA_MODELS=/path/to/models", "OLLAMA_HOST=localhost")
-		err := opt.Customize(&req)
-		require.NoError(t, err)
-		require.Equal(t, "/path/to/models", req.Env["OLLAMA_MODELS"])
-		require.Equal(t, "localhost", req.Env["OLLAMA_HOST"])
+		require.Empty(t, req.Env)
 	})
 
-	t.Run("keyVal/valid/multiple-equals", func(t *testing.T) {
-		opt := ollama.WithUseLocal("OLLAMA_MODELS=/path/to/models", "OLLAMA_HOST=localhost=127.0.0.1")
+	t.Run("valid", func(t *testing.T) {
+		opt := ollama.WithUseLocal(map[string]string{"OLLAMA_MODELS": "/path/to/models", "OLLAMA_HOST": "localhost:1234"})
 		err := opt.Customize(&req)
 		require.NoError(t, err)
 		require.Equal(t, "/path/to/models", req.Env["OLLAMA_MODELS"])
-		require.Equal(t, "localhost=127.0.0.1", req.Env["OLLAMA_HOST"])
-	})
-
-	t.Run("keyVal/invalid/multiple", func(t *testing.T) {
-		opt := ollama.WithUseLocal("OLLAMA_MODELS=/path/to/models", "OLLAMA_HOST")
-		err := opt.Customize(&req)
-		require.Error(t, err)
+		require.Equal(t, "localhost:1234", req.Env["OLLAMA_HOST"])
 	})
 }