|
| 1 | +package latest |
| 2 | + |
| 3 | +import ( |
| 4 | + "testing" |
| 5 | + |
| 6 | + "github.com/goccy/go-yaml" |
| 7 | + "github.com/stretchr/testify/assert" |
| 8 | + "github.com/stretchr/testify/require" |
| 9 | +) |
| 10 | + |
| 11 | +func TestCostConfigValidate(t *testing.T) { |
| 12 | + t.Parallel() |
| 13 | + |
| 14 | + tests := []struct { |
| 15 | + name string |
| 16 | + cost *CostConfig |
| 17 | + wantErr string |
| 18 | + }{ |
| 19 | + {name: "nil is valid"}, |
| 20 | + {name: "all zero means priced free", cost: &CostConfig{}}, |
| 21 | + {name: "positive prices", cost: &CostConfig{Input: 2.5, Output: 10, CacheRead: 0.25, CacheWrite: 3.125}}, |
| 22 | + {name: "negative input", cost: &CostConfig{Input: -1}, wantErr: "cost.input must not be negative, got -1"}, |
| 23 | + {name: "negative output", cost: &CostConfig{Output: -0.5}, wantErr: "cost.output must not be negative, got -0.5"}, |
| 24 | + {name: "negative cache_read", cost: &CostConfig{CacheRead: -2}, wantErr: "cost.cache_read must not be negative, got -2"}, |
| 25 | + {name: "negative cache_write", cost: &CostConfig{CacheWrite: -0.1}, wantErr: "cost.cache_write must not be negative, got -0.1"}, |
| 26 | + } |
| 27 | + |
| 28 | + for _, tt := range tests { |
| 29 | + t.Run(tt.name, func(t *testing.T) { |
| 30 | + t.Parallel() |
| 31 | + err := tt.cost.validate() |
| 32 | + if tt.wantErr != "" { |
| 33 | + require.ErrorContains(t, err, tt.wantErr) |
| 34 | + return |
| 35 | + } |
| 36 | + require.NoError(t, err) |
| 37 | + }) |
| 38 | + } |
| 39 | +} |
| 40 | + |
| 41 | +func TestModelConfigCostYAMLRoundTrip(t *testing.T) { |
| 42 | + t.Parallel() |
| 43 | + |
| 44 | + const in = `provider: internal-llm |
| 45 | +model: gpt-4o |
| 46 | +cost: |
| 47 | + input: 1.25 |
| 48 | + output: 5 |
| 49 | + cache_read: 0.125 |
| 50 | + cache_write: 1.5625 |
| 51 | +` |
| 52 | + var f FlexibleModelConfig |
| 53 | + require.NoError(t, yaml.Unmarshal([]byte(in), &f)) |
| 54 | + |
| 55 | + require.NotNil(t, f.Cost, "cost should be parsed") |
| 56 | + assert.InEpsilon(t, 1.25, f.Cost.Input, 0.0001) |
| 57 | + assert.InEpsilon(t, 5.0, f.Cost.Output, 0.0001) |
| 58 | + assert.InEpsilon(t, 0.125, f.Cost.CacheRead, 0.0001) |
| 59 | + assert.InEpsilon(t, 1.5625, f.Cost.CacheWrite, 0.0001) |
| 60 | + |
| 61 | + // A model carrying a cost override must not collapse to the |
| 62 | + // "provider/model" shorthand on marshal, or the override would be lost. |
| 63 | + assert.False(t, f.isShorthandOnly(), "cost override must defeat shorthand marshalling") |
| 64 | + |
| 65 | + out, err := yaml.Marshal(f) |
| 66 | + require.NoError(t, err) |
| 67 | + |
| 68 | + var rt FlexibleModelConfig |
| 69 | + require.NoError(t, yaml.Unmarshal(out, &rt)) |
| 70 | + require.NotNil(t, rt.Cost, "cost should survive a marshal round-trip; got:\n%s", out) |
| 71 | + assert.InEpsilon(t, 1.25, rt.Cost.Input, 0.0001) |
| 72 | + assert.InEpsilon(t, 5.0, rt.Cost.Output, 0.0001) |
| 73 | +} |
| 74 | + |
| 75 | +func TestConfigValidateRejectsNegativeCost(t *testing.T) { |
| 76 | + t.Parallel() |
| 77 | + |
| 78 | + cfg := Config{ |
| 79 | + Models: map[string]ModelConfig{ |
| 80 | + "main": {Provider: "openai", Model: "gpt-4o", Cost: &CostConfig{Input: -1}}, |
| 81 | + }, |
| 82 | + Agents: Agents{ |
| 83 | + {Name: "root", Model: "main"}, |
| 84 | + }, |
| 85 | + } |
| 86 | + require.ErrorContains(t, cfg.Validate(), "models.main: cost.input must not be negative") |
| 87 | +} |
0 commit comments