ollama/x/cmd/run_test.go

package cmd

import (
	"testing"
)

func TestIsLocalModel(t *testing.T) {
	tests := []struct {
		name      string
		modelName string
		expected  bool
	}{
		{
			name:      "local model without suffix",
			modelName: "llama3.2",
			expected:  true,
		},
		{
			name:      "local model with version",
			modelName: "qwen2.5:7b",
			expected:  true,
		},
		{
			name:      "cloud model",
			modelName: "gpt-4-cloud",
			expected:  false,
		},
		{
			name:      "cloud model with version",
			modelName: "claude-3-cloud",
			expected:  false,
		},
		{
			name:      "empty model name",
			modelName: "",
			expected:  true,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			result := isLocalModel(tt.modelName)
			if result != tt.expected {
				t.Errorf("isLocalModel(%q) = %v, expected %v", tt.modelName, result, tt.expected)
			}
		})
	}
}

func TestIsLocalServer(t *testing.T) {
	tests := []struct {
		name     string
		host     string
		expected bool
	}{
		{
			name:     "empty host (default)",
			host:     "",
			expected: true,
		},
		{
			name:     "localhost",
			host:     "http://localhost:11434",
			expected: true,
		},
		{
			name:     "127.0.0.1",
			host:     "http://127.0.0.1:11434",
			expected: true,
		},
		{
			name:     "custom port on localhost",
			host:     "http://localhost:8080",
			expected: true, // localhost is always considered local
		},
		{
			name:     "remote host",
			host:     "http://ollama.example.com:11434",
			expected: true, // has :11434
		},
		{
			name:     "remote host different port",
			host:     "http://ollama.example.com:8080",
			expected: false,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Setenv("OLLAMA_HOST", tt.host)
			result := isLocalServer()
			if result != tt.expected {
				t.Errorf("isLocalServer() with OLLAMA_HOST=%q = %v, expected %v", tt.host, result, tt.expected)
			}
		})
	}
}

func TestTruncateToolOutput(t *testing.T) {
	// Create outputs of different sizes
	localLimitOutput := make([]byte, 20000)   // > 4k tokens (16k chars)
	defaultLimitOutput := make([]byte, 50000) // > 10k tokens (40k chars)
	for i := range localLimitOutput {
		localLimitOutput[i] = 'a'
	}
	for i := range defaultLimitOutput {
		defaultLimitOutput[i] = 'b'
	}

	tests := []struct {
		name          string
		output        string
		modelName     string
		host          string
		shouldTrim    bool
		expectedLimit int
	}{
		{
			name:          "short output local model",
			output:        "hello world",
			modelName:     "llama3.2",
			host:          "",
			shouldTrim:    false,
			expectedLimit: localModelTokenLimit,
		},
		{
			name:          "long output local model - trimmed at 4k",
			output:        string(localLimitOutput),
			modelName:     "llama3.2",
			host:          "",
			shouldTrim:    true,
			expectedLimit: localModelTokenLimit,
		},
		{
			name:          "long output cloud model - uses 10k limit",
			output:        string(localLimitOutput), // 20k chars, under 10k token limit
			modelName:     "gpt-4-cloud",
			host:          "",
			shouldTrim:    false,
			expectedLimit: defaultTokenLimit,
		},
		{
			name:          "very long output cloud model - trimmed at 10k",
			output:        string(defaultLimitOutput),
			modelName:     "gpt-4-cloud",
			host:          "",
			shouldTrim:    true,
			expectedLimit: defaultTokenLimit,
		},
		{
			name:          "long output remote server - uses 10k limit",
			output:        string(localLimitOutput),
			modelName:     "llama3.2",
			host:          "http://remote.example.com:8080",
			shouldTrim:    false,
			expectedLimit: defaultTokenLimit,
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Setenv("OLLAMA_HOST", tt.host)
			result := truncateToolOutput(tt.output, tt.modelName)

			if tt.shouldTrim {
				maxLen := tt.expectedLimit * charsPerToken
				if len(result) > maxLen+50 { // +50 for the truncation message
					t.Errorf("expected output to be truncated to ~%d chars, got %d", maxLen, len(result))
				}
				if result == tt.output {
					t.Error("expected output to be truncated but it wasn't")
				}
			} else {
				if result != tt.output {
					t.Error("expected output to not be truncated")
				}
			}
		})
	}
}