-
Notifications
You must be signed in to change notification settings - Fork 2
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Support multiline cmds in YAML configuration
Add support for multiline `cmd` configurations allowing for nicer looking configuration YAML files.
- Loading branch information
1 parent
6cf0962
commit be82d1a
Showing
6 changed files
with
185 additions
and
18 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,126 @@ | ||
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) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters