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) }