be82d1a6a0
Add support for multiline `cmd` configurations allowing for nicer looking configuration YAML files.
127 lines
3.2 KiB
Go
127 lines
3.2 KiB
Go
package proxy
|
|
|
|
import (
|
|
"os"
|
|
"path/filepath"
|
|
"testing"
|
|
|
|
"github.com/stretchr/testify/assert"
|
|
)
|
|
|
|
func TestLoadConfig(t *testing.T) {
|
|
// Create a temporary YAML file for testing
|
|
tempDir, err := os.MkdirTemp("", "test-config")
|
|
if err != nil {
|
|
t.Fatalf("Failed to create temporary directory: %v", err)
|
|
}
|
|
defer os.RemoveAll(tempDir)
|
|
|
|
tempFile := filepath.Join(tempDir, "config.yaml")
|
|
content := `models:
|
|
model1:
|
|
cmd: path/to/cmd --arg1 one
|
|
proxy: "http://localhost:8080"
|
|
aliases:
|
|
- "m1"
|
|
- "model-one"
|
|
env:
|
|
- "VAR1=value1"
|
|
- "VAR2=value2"
|
|
checkEndpoint: "/health"
|
|
healthCheckTimeout: 15
|
|
`
|
|
|
|
if err := os.WriteFile(tempFile, []byte(content), 0644); err != nil {
|
|
t.Fatalf("Failed to write temporary file: %v", err)
|
|
}
|
|
|
|
// Load the config and verify
|
|
config, err := LoadConfig(tempFile)
|
|
if err != nil {
|
|
t.Fatalf("Failed to load config: %v", err)
|
|
}
|
|
|
|
expected := &Config{
|
|
Models: map[string]ModelConfig{
|
|
"model1": {
|
|
Cmd: "path/to/cmd --arg1 one",
|
|
Proxy: "http://localhost:8080",
|
|
Aliases: []string{"m1", "model-one"},
|
|
Env: []string{"VAR1=value1", "VAR2=value2"},
|
|
CheckEndpoint: "/health",
|
|
},
|
|
},
|
|
HealthCheckTimeout: 15,
|
|
}
|
|
|
|
assert.Equal(t, expected, config)
|
|
}
|
|
|
|
func TestModelConfigSanitizedCommand(t *testing.T) {
|
|
config := &ModelConfig{
|
|
Cmd: `python model1.py \
|
|
--arg1 value1 \
|
|
--arg2 value2`,
|
|
}
|
|
|
|
args, err := config.SanitizedCommand()
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, []string{"python", "model1.py", "--arg1", "value1", "--arg2", "value2"}, args)
|
|
}
|
|
|
|
func TestFindConfig(t *testing.T) {
|
|
config := &Config{
|
|
Models: map[string]ModelConfig{
|
|
"model1": {
|
|
Cmd: "python model1.py",
|
|
Proxy: "http://localhost:8080",
|
|
Aliases: []string{"m1", "model-one"},
|
|
Env: []string{"VAR1=value1", "VAR2=value2"},
|
|
CheckEndpoint: "/health",
|
|
},
|
|
"model2": {
|
|
Cmd: "python model2.py",
|
|
Proxy: "http://localhost:8081",
|
|
Aliases: []string{"m2", "model-two"},
|
|
Env: []string{"VAR3=value3", "VAR4=value4"},
|
|
CheckEndpoint: "/status",
|
|
},
|
|
},
|
|
HealthCheckTimeout: 10,
|
|
}
|
|
|
|
// Test finding a model by its name
|
|
modelConfig, found := config.FindConfig("model1")
|
|
assert.True(t, found)
|
|
assert.Equal(t, config.Models["model1"], modelConfig)
|
|
|
|
// Test finding a model by its alias
|
|
modelConfig, found = config.FindConfig("m1")
|
|
assert.True(t, found)
|
|
assert.Equal(t, config.Models["model1"], modelConfig)
|
|
|
|
// Test finding a model that does not exist
|
|
modelConfig, found = config.FindConfig("model3")
|
|
assert.False(t, found)
|
|
assert.Equal(t, ModelConfig{}, modelConfig)
|
|
}
|
|
|
|
func TestSanitizeCommand(t *testing.T) {
|
|
// Test a simple command
|
|
args, err := SanitizeCommand("python model1.py")
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, []string{"python", "model1.py"}, args)
|
|
|
|
// Test a command with spaces and newlines
|
|
args, err = SanitizeCommand(`python model1.py \
|
|
--arg1 value1 \
|
|
--arg2 value2`)
|
|
assert.NoError(t, err)
|
|
assert.Equal(t, []string{"python", "model1.py", "--arg1", "value1", "--arg2", "value2"}, args)
|
|
|
|
// Test an empty command
|
|
args, err = SanitizeCommand("")
|
|
assert.Error(t, err)
|
|
assert.Nil(t, args)
|
|
}
|