diff --git a/modules/ollama/examples_test.go b/modules/ollama/examples_test.go index 188be45bbb..3601e0b120 100644 --- 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 index 8dabf5e295..3b95500b72 100644 --- 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 index ed34326d71..4653b65169 100644 --- 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,13 +42,13 @@ 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} } @@ -58,15 +56,9 @@ func WithUseLocal(keyVal ...string) UseLocal { // 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 index f842d15a17..46872d0dd4 100644 --- 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"]) }) }