config_test.go 2.11 KB
Newer Older
1
2
3
package envconfig

import (
4
	"math"
5
	"testing"
6
	"time"
7
8
9
10
11

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

func TestConfig(t *testing.T) {
Daniel Hiltgen's avatar
Daniel Hiltgen committed
12
	t.Setenv("OLLAMA_DEBUG", "")
Michael Yang's avatar
Michael Yang committed
13
14
	require.False(t, Debug())

Daniel Hiltgen's avatar
Daniel Hiltgen committed
15
	t.Setenv("OLLAMA_DEBUG", "false")
Michael Yang's avatar
Michael Yang committed
16
17
	require.False(t, Debug())

Daniel Hiltgen's avatar
Daniel Hiltgen committed
18
	t.Setenv("OLLAMA_DEBUG", "1")
Michael Yang's avatar
Michael Yang committed
19
20
	require.True(t, Debug())

21
22
23
	t.Setenv("OLLAMA_FLASH_ATTENTION", "1")
	LoadConfig()
	require.True(t, FlashAttention)
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
	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)
39
}
40
41

func TestClientFromEnvironment(t *testing.T) {
Michael Yang's avatar
host  
Michael Yang committed
42
	cases := map[string]struct {
43
44
		value  string
		expect string
Michael Yang's avatar
host  
Michael Yang committed
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
	}{
		"empty":               {"", "127.0.0.1:11434"},
		"only address":        {"1.2.3.4", "1.2.3.4:11434"},
		"only port":           {":1234", ":1234"},
		"address and port":    {"1.2.3.4:1234", "1.2.3.4:1234"},
		"hostname":            {"example.com", "example.com:11434"},
		"hostname and port":   {"example.com:1234", "example.com:1234"},
		"zero port":           {":0", ":0"},
		"too large port":      {":66000", ":11434"},
		"too small port":      {":-1", ":11434"},
		"ipv6 localhost":      {"[::1]", "[::1]:11434"},
		"ipv6 world open":     {"[::]", "[::]:11434"},
		"ipv6 no brackets":    {"::1", "[::1]:11434"},
		"ipv6 + port":         {"[::1]:1337", "[::1]:1337"},
		"extra space":         {" 1.2.3.4 ", "1.2.3.4:11434"},
		"extra quotes":        {"\"1.2.3.4\"", "1.2.3.4:11434"},
		"extra space+quotes":  {" \" 1.2.3.4 \" ", "1.2.3.4:11434"},
		"extra single quotes": {"'1.2.3.4'", "1.2.3.4:11434"},
63
64
	}

Michael Yang's avatar
host  
Michael Yang committed
65
66
67
68
69
	for name, tt := range cases {
		t.Run(name, func(t *testing.T) {
			t.Setenv("OLLAMA_HOST", tt.value)
			if host := Host(); host.Host != tt.expect {
				t.Errorf("%s: expected %s, got %s", name, tt.expect, host.Host)
70
71
72
73
			}
		})
	}
}