From 407231c839eba955432e4af00598c3ea33321377 Mon Sep 17 00:00:00 2001 From: Benson Wong Date: Sun, 15 Jun 2025 10:47:36 -0700 Subject: [PATCH 1/9] Refactor all default config values into config.go move all default values for global, model and group configurations so it is all in one place. --- proxy/config.go | 58 +++++++++++++++++++++++++----------- proxy/config_posix_test.go | 42 ++++++++++++++++++++++++++ proxy/config_test.go | 7 +++-- proxy/config_windows_test.go | 40 +++++++++++++++++++++++++ proxy/process.go | 11 ------- 5 files changed, 128 insertions(+), 30 deletions(-) diff --git a/proxy/config.go b/proxy/config.go index bef94e2..22e7ba6 100644 --- a/proxy/config.go +++ b/proxy/config.go @@ -31,6 +31,34 @@ type ModelConfig struct { ConcurrencyLimit int `yaml:"concurrencyLimit"` } +func (m *ModelConfig) UnmarshalYAML(unmarshal func(interface{}) error) error { + type rawModelConfig ModelConfig + defaults := rawModelConfig{ + Cmd: "", + CmdStop: "", + Proxy: "http://localhost:${PORT}", + Aliases: []string{}, + Env: []string{}, + CheckEndpoint: "/health", + UnloadAfter: 0, + Unlisted: false, + UseModelName: "", + ConcurrencyLimit: 0, + } + + // the default cmdStop to taskkill /f /t /pid ${PID} + if runtime.GOOS == "windows" { + defaults.CmdStop = "taskkill /f /t /pid ${PID}" + } + + if err := unmarshal(&defaults); err != nil { + return err + } + + *m = ModelConfig(defaults) + return nil +} + func (m *ModelConfig) SanitizedCommand() ([]string, error) { return SanitizeCommand(m.Cmd) } @@ -111,26 +139,23 @@ func LoadConfigFromReader(r io.Reader) (Config, error) { return Config{}, err } - var config Config + // default configuration values + config := Config{ + HealthCheckTimeout: 120, + StartPort: 5800, + LogLevel: "info", + } err = yaml.Unmarshal(data, &config) if err != nil { return Config{}, err } - if config.HealthCheckTimeout == 0 { - // this high default timeout helps avoid failing health checks - // for configurations that wait for docker or have slower startup - config.HealthCheckTimeout = 120 - } else if config.HealthCheckTimeout < 15 { + if config.HealthCheckTimeout < 15 { // set a minimum of 15 seconds config.HealthCheckTimeout = 15 } - // set default port ranges - if config.StartPort == 0 { - // default to 5800 - config.StartPort = 5800 - } else if config.StartPort < 1 { + if config.StartPort < 1 { return Config{}, fmt.Errorf("startPort must be greater than 1") } @@ -180,6 +205,11 @@ func LoadConfigFromReader(r io.Reader) (Config, error) { for _, modelId := range modelIds { modelConfig := config.Models[modelId] + // enforce ${PORT} used in both cmd and proxy + if !strings.Contains(modelConfig.Cmd, "${PORT}") && strings.Contains(modelConfig.Proxy, "${PORT}") { + return Config{}, fmt.Errorf("model %s requires a proxy value when not using automatic ${PORT}", modelId) + } + // go through model config fields: cmd, cmdStop, proxy, checkEndPoint and replace macros with macro values for macroName, macroValue := range config.Macros { macroSlug := fmt.Sprintf("${%s}", macroName) @@ -191,17 +221,11 @@ func LoadConfigFromReader(r io.Reader) (Config, error) { // only iterate over models that use ${PORT} to keep port numbers from increasing unnecessarily if strings.Contains(modelConfig.Cmd, "${PORT}") || strings.Contains(modelConfig.Proxy, "${PORT}") || strings.Contains(modelConfig.CmdStop, "${PORT}") { - if modelConfig.Proxy == "" { - modelConfig.Proxy = "http://localhost:${PORT}" - } - nextPortStr := strconv.Itoa(nextPort) modelConfig.Cmd = strings.ReplaceAll(modelConfig.Cmd, "${PORT}", nextPortStr) modelConfig.CmdStop = strings.ReplaceAll(modelConfig.CmdStop, "${PORT}", nextPortStr) modelConfig.Proxy = strings.ReplaceAll(modelConfig.Proxy, "${PORT}", nextPortStr) nextPort++ - } else if modelConfig.Proxy == "" { - return Config{}, fmt.Errorf("model %s requires a proxy value when not using automatic ${PORT}", modelId) } // make sure there are no unknown macros that have not been replaced diff --git a/proxy/config_posix_test.go b/proxy/config_posix_test.go index 6f9c452..8619300 100644 --- a/proxy/config_posix_test.go +++ b/proxy/config_posix_test.go @@ -3,6 +3,7 @@ package proxy import ( + "strings" "testing" "github.com/stretchr/testify/assert" @@ -40,3 +41,44 @@ func TestConfig_SanitizeCommand(t *testing.T) { assert.Error(t, err) assert.Nil(t, args) } + +// Test the default values are automatically set for global, model and group configurations +// after loading the configuration +func TestConfig_DefaultValues(t *testing.T) { + content := ` +models: + model1: + cmd: path/to/cmd --port ${PORT} +` + + config, err := LoadConfigFromReader(strings.NewReader(content)) + assert.NoError(t, err) + assert.Equal(t, 120, config.HealthCheckTimeout) + assert.Equal(t, 5800, config.StartPort) + assert.Equal(t, "info", config.LogLevel) + + // Test default group exists + defaultGroup, exists := config.Groups["(default)"] + assert.True(t, exists, "default group should exist") + if assert.NotNil(t, defaultGroup, "default group should not be nil") { + assert.Equal(t, true, defaultGroup.Swap) + assert.Equal(t, true, defaultGroup.Exclusive) + assert.Equal(t, false, defaultGroup.Persistent) + assert.Equal(t, []string{"model1"}, defaultGroup.Members) + } + + model1, exists := config.Models["model1"] + assert.True(t, exists, "model1 should exist") + if assert.NotNil(t, model1, "model1 should not be nil") { + assert.Equal(t, "path/to/cmd --port 5800", model1.Cmd) // has the port replaced + assert.Equal(t, "", model1.CmdStop) + assert.Equal(t, "http://localhost:5800", model1.Proxy) + assert.Equal(t, "/health", model1.CheckEndpoint) + assert.Equal(t, []string{}, model1.Aliases) + assert.Equal(t, []string{}, model1.Env) + assert.Equal(t, 0, model1.UnloadAfter) + assert.Equal(t, false, model1.Unlisted) + assert.Equal(t, "", model1.UseModelName) + assert.Equal(t, 0, model1.ConcurrencyLimit) + } +} diff --git a/proxy/config_test.go b/proxy/config_test.go index 5ce93e0..eabcd65 100644 --- a/proxy/config_test.go +++ b/proxy/config_test.go @@ -77,6 +77,7 @@ groups: } expected := Config{ + LogLevel: "info", StartPort: 5800, Macros: map[string]string{ "svr-path": "path/to/server", @@ -93,20 +94,22 @@ groups: Cmd: "path/to/server --arg1 one", Proxy: "http://localhost:8081", Aliases: []string{"m2"}, - Env: nil, + Env: []string{}, CheckEndpoint: "/", }, "model3": { Cmd: "path/to/cmd --arg1 one", Proxy: "http://localhost:8081", Aliases: []string{"mthree"}, - Env: nil, + Env: []string{}, CheckEndpoint: "/", }, "model4": { Cmd: "path/to/cmd --arg1 one", Proxy: "http://localhost:8082", CheckEndpoint: "/", + Aliases: []string{}, + Env: []string{}, }, }, HealthCheckTimeout: 15, diff --git a/proxy/config_windows_test.go b/proxy/config_windows_test.go index d301a06..fe0f589 100644 --- a/proxy/config_windows_test.go +++ b/proxy/config_windows_test.go @@ -3,6 +3,7 @@ package proxy import ( + "strings" "testing" "github.com/stretchr/testify/assert" @@ -39,3 +40,42 @@ func TestConfig_SanitizeCommand(t *testing.T) { assert.Error(t, err) assert.Nil(t, args) } + +func TestConfig_DefaultValuesWindows(t *testing.T) { + content := ` +models: + model1: + cmd: path/to/cmd --port ${PORT} +` + + config, err := LoadConfigFromReader(strings.NewReader(content)) + assert.NoError(t, err) + assert.Equal(t, 120, config.HealthCheckTimeout) + assert.Equal(t, 5800, config.StartPort) + assert.Equal(t, "info", config.LogLevel) + + // Test default group exists + defaultGroup, exists := config.Groups["(default)"] + assert.True(t, exists, "default group should exist") + if assert.NotNil(t, defaultGroup, "default group should not be nil") { + assert.Equal(t, true, defaultGroup.Swap) + assert.Equal(t, true, defaultGroup.Exclusive) + assert.Equal(t, false, defaultGroup.Persistent) + assert.Equal(t, []string{"model1"}, defaultGroup.Members) + } + + model1, exists := config.Models["model1"] + assert.True(t, exists, "model1 should exist") + if assert.NotNil(t, model1, "model1 should not be nil") { + assert.Equal(t, "path/to/cmd --port 5800", model1.Cmd) // has the port replaced + assert.Equal(t, "", model1.CmdStop) + assert.Equal(t, "http://localhost:5800", model1.Proxy) + assert.Equal(t, "/health", model1.CheckEndpoint) + assert.Equal(t, []string{}, model1.Aliases) + assert.Equal(t, []string{}, model1.Env) + assert.Equal(t, 0, model1.UnloadAfter) + assert.Equal(t, false, model1.Unlisted) + assert.Equal(t, "", model1.UseModelName) + assert.Equal(t, 0, model1.ConcurrencyLimit) + } +} diff --git a/proxy/process.go b/proxy/process.go index 6ffc18c..487efcc 100644 --- a/proxy/process.go +++ b/proxy/process.go @@ -8,7 +8,6 @@ import ( "net/http" "net/url" "os/exec" - "runtime" "strconv" "strings" "sync" @@ -232,11 +231,6 @@ func (p *Process) start() error { // a "none" means don't check for health ... I could have picked a better word :facepalm: if checkEndpoint != "none" { - // keep default behaviour - if checkEndpoint == "" { - checkEndpoint = "/health" - } - proxyTo := p.config.Proxy healthURL, err := url.JoinPath(proxyTo, checkEndpoint) if err != nil { @@ -523,11 +517,6 @@ func (p *Process) cmdStopUpstreamProcess() error { return fmt.Errorf("<%s> process is nil or cmd is nil, skipping graceful stop", p.ID) } - // the default cmdStop to taskkill /f /t /pid ${PID} - if runtime.GOOS == "windows" && strings.TrimSpace(p.config.CmdStop) == "" { - p.config.CmdStop = "taskkill /f /t /pid ${PID}" - } - if p.config.CmdStop != "" { // replace ${PID} with the pid of the process stopArgs, err := SanitizeCommand(strings.ReplaceAll(p.config.CmdStop, "${PID}", fmt.Sprintf("%d", p.cmd.Process.Pid))) From d9c0b5dc1ad29c6dbf5037e8f94f2167843df444 Mon Sep 17 00:00:00 2001 From: Benson Wong Date: Sun, 15 Jun 2025 11:48:21 -0700 Subject: [PATCH 2/9] . --- proxy/config.go | 2 +- proxy/config_posix_test.go | 2 +- proxy/config_test.go | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/proxy/config.go b/proxy/config.go index 22e7ba6..4faa9d3 100644 --- a/proxy/config.go +++ b/proxy/config.go @@ -207,7 +207,7 @@ func LoadConfigFromReader(r io.Reader) (Config, error) { // enforce ${PORT} used in both cmd and proxy if !strings.Contains(modelConfig.Cmd, "${PORT}") && strings.Contains(modelConfig.Proxy, "${PORT}") { - return Config{}, fmt.Errorf("model %s requires a proxy value when not using automatic ${PORT}", modelId) + return Config{}, fmt.Errorf("model %s: proxy uses ${PORT} but cmd does not - ${PORT} is only available when used in cmd", modelId) } // go through model config fields: cmd, cmdStop, proxy, checkEndPoint and replace macros with macro values diff --git a/proxy/config_posix_test.go b/proxy/config_posix_test.go index 8619300..b1504d5 100644 --- a/proxy/config_posix_test.go +++ b/proxy/config_posix_test.go @@ -44,7 +44,7 @@ func TestConfig_SanitizeCommand(t *testing.T) { // Test the default values are automatically set for global, model and group configurations // after loading the configuration -func TestConfig_DefaultValues(t *testing.T) { +func TestConfig_DefaultValuesPosix(t *testing.T) { content := ` models: model1: diff --git a/proxy/config_test.go b/proxy/config_test.go index eabcd65..4a30a00 100644 --- a/proxy/config_test.go +++ b/proxy/config_test.go @@ -336,7 +336,7 @@ models: cmd: svr --port 111 ` _, err := LoadConfigFromReader(strings.NewReader(content)) - assert.Equal(t, "model model1 requires a proxy value when not using automatic ${PORT}", err.Error()) + assert.Equal(t, "model model1: proxy uses ${PORT} but cmd does not - ${PORT} is only available when used in cmd", err.Error()) }) } From b5146ab0c6dfa0d092f3434a1c670ea564d8d198 Mon Sep 17 00:00:00 2001 From: Benson Wong Date: Sun, 15 Jun 2025 11:59:41 -0700 Subject: [PATCH 3/9] allow ${PID} in cmdStop to pass validation checks --- proxy/config.go | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/proxy/config.go b/proxy/config.go index 4faa9d3..ab709b0 100644 --- a/proxy/config.go +++ b/proxy/config.go @@ -35,7 +35,7 @@ func (m *ModelConfig) UnmarshalYAML(unmarshal func(interface{}) error) error { type rawModelConfig ModelConfig defaults := rawModelConfig{ Cmd: "", - CmdStop: "", + CmdStop: "test ${PID}", Proxy: "http://localhost:${PORT}", Aliases: []string{}, Env: []string{}, @@ -241,6 +241,9 @@ func LoadConfigFromReader(r io.Reader) (Config, error) { matches := macroPattern.FindAllStringSubmatch(fieldValue, -1) for _, match := range matches { macroName := match[1] + if macroName == "PID" && fieldName == "cmdStop" { + continue // this is ok, has to be replaced by process later + } if _, exists := config.Macros[macroName]; !exists { return Config{}, fmt.Errorf("unknown macro '${%s}' found in %s.%s", macroName, modelId, fieldName) } From 3d7873339e457f6abdf5ee78f4ccb747a640f248 Mon Sep 17 00:00:00 2001 From: Benson Wong Date: Sun, 15 Jun 2025 12:00:38 -0700 Subject: [PATCH 4/9] fix wrong cmdStop windows default config --- proxy/config_windows_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/proxy/config_windows_test.go b/proxy/config_windows_test.go index fe0f589..94f1410 100644 --- a/proxy/config_windows_test.go +++ b/proxy/config_windows_test.go @@ -68,7 +68,7 @@ models: assert.True(t, exists, "model1 should exist") if assert.NotNil(t, model1, "model1 should not be nil") { assert.Equal(t, "path/to/cmd --port 5800", model1.Cmd) // has the port replaced - assert.Equal(t, "", model1.CmdStop) + assert.Equal(t, "taskkill /f /t /pid ${PID}", model1.CmdStop) assert.Equal(t, "http://localhost:5800", model1.Proxy) assert.Equal(t, "/health", model1.CheckEndpoint) assert.Equal(t, []string{}, model1.Aliases) From 47283ac67f6c30e7ac8e7ffe955f54cb7fd24692 Mon Sep 17 00:00:00 2001 From: Benson Wong Date: Sun, 15 Jun 2025 12:02:29 -0700 Subject: [PATCH 5/9] remove debug test value --- proxy/config.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/proxy/config.go b/proxy/config.go index ab709b0..9d564f6 100644 --- a/proxy/config.go +++ b/proxy/config.go @@ -35,7 +35,7 @@ func (m *ModelConfig) UnmarshalYAML(unmarshal func(interface{}) error) error { type rawModelConfig ModelConfig defaults := rawModelConfig{ Cmd: "", - CmdStop: "test ${PID}", + CmdStop: "", Proxy: "http://localhost:${PORT}", Aliases: []string{}, Env: []string{}, From 13d2d86ae219f62cbb8fc6ea1bd782d950c77064 Mon Sep 17 00:00:00 2001 From: Benson Wong Date: Sun, 15 Jun 2025 12:13:35 -0700 Subject: [PATCH 6/9] use defaults from unmarshalling yaml in tests --- proxy/helpers_test.go | 17 ++++++++++++----- 1 file changed, 12 insertions(+), 5 deletions(-) diff --git a/proxy/helpers_test.go b/proxy/helpers_test.go index 0cf7c03..413d63e 100644 --- a/proxy/helpers_test.go +++ b/proxy/helpers_test.go @@ -9,6 +9,7 @@ import ( "testing" "github.com/gin-gonic/gin" + "gopkg.in/yaml.v3" ) var ( @@ -70,10 +71,16 @@ func getTestSimpleResponderConfig(expectedMessage string) ModelConfig { func getTestSimpleResponderConfigPort(expectedMessage string, port int) ModelConfig { binaryPath := getSimpleResponderPath() - // Create a process configuration - return ModelConfig{ - Cmd: fmt.Sprintf("%s --port %d --silent --respond %s", binaryPath, port, expectedMessage), - Proxy: fmt.Sprintf("http://127.0.0.1:%d", port), - CheckEndpoint: "/health", + // Create a YAML string with just the values we want to set + yamlStr := fmt.Sprintf(` +cmd: "%s --port %d --silent --respond %s" +proxy: "http://127.0.0.1:%d" +`, binaryPath, port, expectedMessage, port) + + var cfg ModelConfig + if err := yaml.Unmarshal([]byte(yamlStr), &cfg); err != nil { + panic(fmt.Sprintf("failed to unmarshal test config: %v", err)) } + + return cfg } From 3468d4de04f9fb879ec59c6dfe27655a64bb3e9d Mon Sep 17 00:00:00 2001 From: Benson Wong Date: Sun, 15 Jun 2025 12:19:50 -0700 Subject: [PATCH 7/9] . --- proxy/helpers_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/proxy/helpers_test.go b/proxy/helpers_test.go index 413d63e..da9395a 100644 --- a/proxy/helpers_test.go +++ b/proxy/helpers_test.go @@ -73,13 +73,13 @@ func getTestSimpleResponderConfigPort(expectedMessage string, port int) ModelCon // Create a YAML string with just the values we want to set yamlStr := fmt.Sprintf(` -cmd: "%s --port %d --silent --respond %s" +cmd: "%q --port %d --silent --respond %q" proxy: "http://127.0.0.1:%d" `, binaryPath, port, expectedMessage, port) var cfg ModelConfig if err := yaml.Unmarshal([]byte(yamlStr), &cfg); err != nil { - panic(fmt.Sprintf("failed to unmarshal test config: %v", err)) + panic(fmt.Sprintf("failed to unmarshal test config: %v in [%s]", err, yamlStr)) } return cfg From 1e8ed6512f298165425e01d3c8571c2876f93080 Mon Sep 17 00:00:00 2001 From: Benson Wong Date: Sun, 15 Jun 2025 12:20:48 -0700 Subject: [PATCH 8/9] . --- proxy/helpers_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/proxy/helpers_test.go b/proxy/helpers_test.go index da9395a..ccea359 100644 --- a/proxy/helpers_test.go +++ b/proxy/helpers_test.go @@ -73,7 +73,7 @@ func getTestSimpleResponderConfigPort(expectedMessage string, port int) ModelCon // Create a YAML string with just the values we want to set yamlStr := fmt.Sprintf(` -cmd: "%q --port %d --silent --respond %q" +cmd: '%s --port %d --silent --respond %s' proxy: "http://127.0.0.1:%d" `, binaryPath, port, expectedMessage, port) From b9f0d33d5476915c729b41c0e15ca22ffdeffc01 Mon Sep 17 00:00:00 2001 From: Benson Wong Date: Sun, 15 Jun 2025 12:26:06 -0700 Subject: [PATCH 9/9] more fixing cross platform config tests --- proxy/config_posix_test.go | 142 ++++++++++++++++++++++++++++++++++ proxy/config_test.go | 142 ---------------------------------- proxy/config_windows_test.go | 146 +++++++++++++++++++++++++++++++++++ proxy/process_test.go | 1 + 4 files changed, 289 insertions(+), 142 deletions(-) diff --git a/proxy/config_posix_test.go b/proxy/config_posix_test.go index b1504d5..2023eb7 100644 --- a/proxy/config_posix_test.go +++ b/proxy/config_posix_test.go @@ -3,6 +3,8 @@ package proxy import ( + "os" + "path/filepath" "strings" "testing" @@ -82,3 +84,143 @@ models: assert.Equal(t, 0, model1.ConcurrencyLimit) } } + +func TestConfig_LoadPosix(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 := ` +macros: + svr-path: "path/to/server" +models: + model1: + cmd: path/to/cmd --arg1 one + proxy: "http://localhost:8080" + aliases: + - "m1" + - "model-one" + env: + - "VAR1=value1" + - "VAR2=value2" + checkEndpoint: "/health" + model2: + cmd: ${svr-path} --arg1 one + proxy: "http://localhost:8081" + aliases: + - "m2" + checkEndpoint: "/" + model3: + cmd: path/to/cmd --arg1 one + proxy: "http://localhost:8081" + aliases: + - "mthree" + checkEndpoint: "/" + model4: + cmd: path/to/cmd --arg1 one + proxy: "http://localhost:8082" + checkEndpoint: "/" + +healthCheckTimeout: 15 +profiles: + test: + - model1 + - model2 +groups: + group1: + swap: true + exclusive: false + members: ["model2"] + forever: + exclusive: false + persistent: true + members: + - "model4" +` + + 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{ + LogLevel: "info", + StartPort: 5800, + Macros: map[string]string{ + "svr-path": "path/to/server", + }, + 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", + }, + "model2": { + Cmd: "path/to/server --arg1 one", + Proxy: "http://localhost:8081", + Aliases: []string{"m2"}, + Env: []string{}, + CheckEndpoint: "/", + }, + "model3": { + Cmd: "path/to/cmd --arg1 one", + Proxy: "http://localhost:8081", + Aliases: []string{"mthree"}, + Env: []string{}, + CheckEndpoint: "/", + }, + "model4": { + Cmd: "path/to/cmd --arg1 one", + Proxy: "http://localhost:8082", + CheckEndpoint: "/", + Aliases: []string{}, + Env: []string{}, + }, + }, + HealthCheckTimeout: 15, + Profiles: map[string][]string{ + "test": {"model1", "model2"}, + }, + aliases: map[string]string{ + "m1": "model1", + "model-one": "model1", + "m2": "model2", + "mthree": "model3", + }, + Groups: map[string]GroupConfig{ + DEFAULT_GROUP_ID: { + Swap: true, + Exclusive: true, + Members: []string{"model1", "model3"}, + }, + "group1": { + Swap: true, + Exclusive: false, + Members: []string{"model2"}, + }, + "forever": { + Swap: true, + Exclusive: false, + Persistent: true, + Members: []string{"model4"}, + }, + }, + } + + assert.Equal(t, expected, config) + + realname, found := config.RealModelName("m1") + assert.True(t, found) + assert.Equal(t, "model1", realname) +} diff --git a/proxy/config_test.go b/proxy/config_test.go index 4a30a00..49d4711 100644 --- a/proxy/config_test.go +++ b/proxy/config_test.go @@ -1,154 +1,12 @@ package proxy import ( - "os" - "path/filepath" "strings" "testing" "github.com/stretchr/testify/assert" ) -func TestConfig_Load(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 := ` -macros: - svr-path: "path/to/server" -models: - model1: - cmd: path/to/cmd --arg1 one - proxy: "http://localhost:8080" - aliases: - - "m1" - - "model-one" - env: - - "VAR1=value1" - - "VAR2=value2" - checkEndpoint: "/health" - model2: - cmd: ${svr-path} --arg1 one - proxy: "http://localhost:8081" - aliases: - - "m2" - checkEndpoint: "/" - model3: - cmd: path/to/cmd --arg1 one - proxy: "http://localhost:8081" - aliases: - - "mthree" - checkEndpoint: "/" - model4: - cmd: path/to/cmd --arg1 one - proxy: "http://localhost:8082" - checkEndpoint: "/" - -healthCheckTimeout: 15 -profiles: - test: - - model1 - - model2 -groups: - group1: - swap: true - exclusive: false - members: ["model2"] - forever: - exclusive: false - persistent: true - members: - - "model4" -` - - 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{ - LogLevel: "info", - StartPort: 5800, - Macros: map[string]string{ - "svr-path": "path/to/server", - }, - 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", - }, - "model2": { - Cmd: "path/to/server --arg1 one", - Proxy: "http://localhost:8081", - Aliases: []string{"m2"}, - Env: []string{}, - CheckEndpoint: "/", - }, - "model3": { - Cmd: "path/to/cmd --arg1 one", - Proxy: "http://localhost:8081", - Aliases: []string{"mthree"}, - Env: []string{}, - CheckEndpoint: "/", - }, - "model4": { - Cmd: "path/to/cmd --arg1 one", - Proxy: "http://localhost:8082", - CheckEndpoint: "/", - Aliases: []string{}, - Env: []string{}, - }, - }, - HealthCheckTimeout: 15, - Profiles: map[string][]string{ - "test": {"model1", "model2"}, - }, - aliases: map[string]string{ - "m1": "model1", - "model-one": "model1", - "m2": "model2", - "mthree": "model3", - }, - Groups: map[string]GroupConfig{ - DEFAULT_GROUP_ID: { - Swap: true, - Exclusive: true, - Members: []string{"model1", "model3"}, - }, - "group1": { - Swap: true, - Exclusive: false, - Members: []string{"model2"}, - }, - "forever": { - Swap: true, - Exclusive: false, - Persistent: true, - Members: []string{"model4"}, - }, - }, - } - - assert.Equal(t, expected, config) - - realname, found := config.RealModelName("m1") - assert.True(t, found) - assert.Equal(t, "model1", realname) -} - func TestConfig_GroupMemberIsUnique(t *testing.T) { content := ` models: diff --git a/proxy/config_windows_test.go b/proxy/config_windows_test.go index 94f1410..d5cb50c 100644 --- a/proxy/config_windows_test.go +++ b/proxy/config_windows_test.go @@ -3,6 +3,8 @@ package proxy import ( + "os" + "path/filepath" "strings" "testing" @@ -79,3 +81,147 @@ models: assert.Equal(t, 0, model1.ConcurrencyLimit) } } + +func TestConfig_LoadWindows(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 := ` +macros: + svr-path: "path/to/server" +models: + model1: + cmd: path/to/cmd --arg1 one + proxy: "http://localhost:8080" + aliases: + - "m1" + - "model-one" + env: + - "VAR1=value1" + - "VAR2=value2" + checkEndpoint: "/health" + model2: + cmd: ${svr-path} --arg1 one + proxy: "http://localhost:8081" + aliases: + - "m2" + checkEndpoint: "/" + model3: + cmd: path/to/cmd --arg1 one + proxy: "http://localhost:8081" + aliases: + - "mthree" + checkEndpoint: "/" + model4: + cmd: path/to/cmd --arg1 one + proxy: "http://localhost:8082" + checkEndpoint: "/" + +healthCheckTimeout: 15 +profiles: + test: + - model1 + - model2 +groups: + group1: + swap: true + exclusive: false + members: ["model2"] + forever: + exclusive: false + persistent: true + members: + - "model4" +` + + 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{ + LogLevel: "info", + StartPort: 5800, + Macros: map[string]string{ + "svr-path": "path/to/server", + }, + Models: map[string]ModelConfig{ + "model1": { + Cmd: "path/to/cmd --arg1 one", + CmdStop: "taskkill /f /t /pid ${PID}", + Proxy: "http://localhost:8080", + Aliases: []string{"m1", "model-one"}, + Env: []string{"VAR1=value1", "VAR2=value2"}, + CheckEndpoint: "/health", + }, + "model2": { + Cmd: "path/to/server --arg1 one", + CmdStop: "taskkill /f /t /pid ${PID}", + Proxy: "http://localhost:8081", + Aliases: []string{"m2"}, + Env: []string{}, + CheckEndpoint: "/", + }, + "model3": { + Cmd: "path/to/cmd --arg1 one", + CmdStop: "taskkill /f /t /pid ${PID}", + Proxy: "http://localhost:8081", + Aliases: []string{"mthree"}, + Env: []string{}, + CheckEndpoint: "/", + }, + "model4": { + Cmd: "path/to/cmd --arg1 one", + CmdStop: "taskkill /f /t /pid ${PID}", + Proxy: "http://localhost:8082", + CheckEndpoint: "/", + Aliases: []string{}, + Env: []string{}, + }, + }, + HealthCheckTimeout: 15, + Profiles: map[string][]string{ + "test": {"model1", "model2"}, + }, + aliases: map[string]string{ + "m1": "model1", + "model-one": "model1", + "m2": "model2", + "mthree": "model3", + }, + Groups: map[string]GroupConfig{ + DEFAULT_GROUP_ID: { + Swap: true, + Exclusive: true, + Members: []string{"model1", "model3"}, + }, + "group1": { + Swap: true, + Exclusive: false, + Members: []string{"model2"}, + }, + "forever": { + Swap: true, + Exclusive: false, + Persistent: true, + Members: []string{"model4"}, + }, + }, + } + + assert.Equal(t, expected, config) + + realname, found := config.RealModelName("m1") + assert.True(t, found) + assert.Equal(t, "model1", realname) +} diff --git a/proxy/process_test.go b/proxy/process_test.go index f734ac8..a87b4db 100644 --- a/proxy/process_test.go +++ b/proxy/process_test.go @@ -405,6 +405,7 @@ func TestProcess_ForceStopWithKill(t *testing.T) { Cmd: fmt.Sprintf("%s --port %d --respond %s --silent --ignore-sig-term", binaryPath, port, expectedMessage), Proxy: fmt.Sprintf("http://127.0.0.1:%d", port), CheckEndpoint: "/health", + CmdStop: "taskkill /f /t /pid ${PID}", } process := NewProcess("stop_immediate", 2, config, debugLogger, debugLogger)