ollama/envconfig/config_test.go

package envconfig

import (
	"fmt"
	"math"
	"net"
	"testing"
	"time"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestConfig(t *testing.T) {
	t.Setenv("OLLAMA_DEBUG", "")
	require.False(t, Debug())

	t.Setenv("OLLAMA_DEBUG", "false")
	require.False(t, Debug())

	t.Setenv("OLLAMA_DEBUG", "1")
	require.True(t, Debug())

	t.Setenv("OLLAMA_FLASH_ATTENTION", "1")
	LoadConfig()
	require.True(t, FlashAttention)
	t.Setenv("OLLAMA_KEEP_ALIVE", "")
	LoadConfig()
	require.Equal(t, 5*time.Minute, KeepAlive)
	t.Setenv("OLLAMA_KEEP_ALIVE", "3")
	LoadConfig()
	require.Equal(t, 3*time.Second, KeepAlive)
	t.Setenv("OLLAMA_KEEP_ALIVE", "1h")
	LoadConfig()
	require.Equal(t, 1*time.Hour, KeepAlive)
	t.Setenv("OLLAMA_KEEP_ALIVE", "-1s")
	LoadConfig()
	require.Equal(t, time.Duration(math.MaxInt64), KeepAlive)
	t.Setenv("OLLAMA_KEEP_ALIVE", "-1")
	LoadConfig()
	require.Equal(t, time.Duration(math.MaxInt64), KeepAlive)
}

func TestClientFromEnvironment(t *testing.T) {
	type testCase struct {
		value  string
		expect string
		err    error
	}

	hostTestCases := map[string]*testCase{
		"empty":               {value: "", expect: "127.0.0.1:11434"},
		"only address":        {value: "1.2.3.4", expect: "1.2.3.4:11434"},
		"only port":           {value: ":1234", expect: ":1234"},
		"address and port":    {value: "1.2.3.4:1234", expect: "1.2.3.4:1234"},
		"hostname":            {value: "example.com", expect: "example.com:11434"},
		"hostname and port":   {value: "example.com:1234", expect: "example.com:1234"},
		"zero port":           {value: ":0", expect: ":0"},
		"too large port":      {value: ":66000", err: ErrInvalidHostPort},
		"too small port":      {value: ":-1", err: ErrInvalidHostPort},
		"ipv6 localhost":      {value: "[::1]", expect: "[::1]:11434"},
		"ipv6 world open":     {value: "[::]", expect: "[::]:11434"},
		"ipv6 no brackets":    {value: "::1", expect: "[::1]:11434"},
		"ipv6 + port":         {value: "[::1]:1337", expect: "[::1]:1337"},
		"extra space":         {value: " 1.2.3.4 ", expect: "1.2.3.4:11434"},
		"extra quotes":        {value: "\"1.2.3.4\"", expect: "1.2.3.4:11434"},
		"extra space+quotes":  {value: " \" 1.2.3.4 \" ", expect: "1.2.3.4:11434"},
		"extra single quotes": {value: "'1.2.3.4'", expect: "1.2.3.4:11434"},
	}

	for k, v := range hostTestCases {
		t.Run(k, func(t *testing.T) {
			t.Setenv("OLLAMA_HOST", v.value)
			LoadConfig()

			oh, err := getOllamaHost()
			if err != v.err {
				t.Fatalf("expected %s, got %s", v.err, err)
			}

			if err == nil {
				host := net.JoinHostPort(oh.Host, oh.Port)
				assert.Equal(t, v.expect, host, fmt.Sprintf("%s: expected %s, got %s", k, v.expect, host))
			}
		})
	}
}