unkinben
ab3b02a48e
Populate the repo with the Vault/OpenBao dynamic secrets engine that mints LiteLLM virtual keys scoped by model, spending limit, and lease TTL. - Secrets backend: config, roles, creds paths and a revocable litellm_key type - LiteLLM API client (generate/update/delete/info) with master-key auth - Unit tests (mock LiteLLM) and a docker-compose e2e against both Vault and OpenBao proving the same binary works on each - Makefile, woodpecker CI (build/test/pre-commit), pre-commit config
123 lines
3.0 KiB
Go
123 lines
3.0 KiB
Go
package litellm
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
|
|
"github.com/hashicorp/vault/sdk/logical"
|
|
)
|
|
|
|
func TestRole_WriteReadListDelete(t *testing.T) {
|
|
b, s := getTestBackend(t)
|
|
ctx := context.Background()
|
|
|
|
resp, err := b.HandleRequest(ctx, &logical.Request{
|
|
Operation: logical.CreateOperation,
|
|
Path: "roles/team-a",
|
|
Storage: s,
|
|
Data: map[string]interface{}{
|
|
"models": "gpt-4,gpt-3.5-turbo",
|
|
"max_budget": 50.5,
|
|
"key_alias_prefix": "team-a",
|
|
"ttl": "1h",
|
|
"max_ttl": "24h",
|
|
"metadata": "env=prod,team=a",
|
|
},
|
|
})
|
|
if err != nil || (resp != nil && resp.IsError()) {
|
|
t.Fatalf("write role: err=%v resp=%v", err, resp)
|
|
}
|
|
|
|
resp, err = b.HandleRequest(ctx, &logical.Request{
|
|
Operation: logical.ReadOperation,
|
|
Path: "roles/team-a",
|
|
Storage: s,
|
|
})
|
|
if err != nil || resp == nil {
|
|
t.Fatalf("read role: err=%v resp=%v", err, resp)
|
|
}
|
|
models := resp.Data["models"].([]string)
|
|
if len(models) != 2 || models[0] != "gpt-4" {
|
|
t.Fatalf("unexpected models: %v", models)
|
|
}
|
|
if resp.Data["max_budget"].(float64) != 50.5 {
|
|
t.Fatalf("unexpected budget: %v", resp.Data["max_budget"])
|
|
}
|
|
if resp.Data["ttl"].(int64) != 3600 {
|
|
t.Fatalf("unexpected ttl: %v", resp.Data["ttl"])
|
|
}
|
|
if resp.Data["max_ttl"].(int64) != 86400 {
|
|
t.Fatalf("unexpected max_ttl: %v", resp.Data["max_ttl"])
|
|
}
|
|
|
|
// List.
|
|
resp, err = b.HandleRequest(ctx, &logical.Request{
|
|
Operation: logical.ListOperation,
|
|
Path: "roles/",
|
|
Storage: s,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("list roles: %v", err)
|
|
}
|
|
keys := resp.Data["keys"].([]string)
|
|
if len(keys) != 1 || keys[0] != "team-a" {
|
|
t.Fatalf("unexpected role list: %v", keys)
|
|
}
|
|
|
|
// Delete.
|
|
if _, err := b.HandleRequest(ctx, &logical.Request{
|
|
Operation: logical.DeleteOperation,
|
|
Path: "roles/team-a",
|
|
Storage: s,
|
|
}); err != nil {
|
|
t.Fatalf("delete role: %v", err)
|
|
}
|
|
role, _ := b.getRole(ctx, s, "team-a")
|
|
if role != nil {
|
|
t.Fatal("expected role to be gone after delete")
|
|
}
|
|
}
|
|
|
|
func TestRole_TTLGreaterThanMaxTTLRejected(t *testing.T) {
|
|
b, s := getTestBackend(t)
|
|
ctx := context.Background()
|
|
|
|
resp, err := b.HandleRequest(ctx, &logical.Request{
|
|
Operation: logical.CreateOperation,
|
|
Path: "roles/bad",
|
|
Storage: s,
|
|
Data: map[string]interface{}{
|
|
"ttl": "48h",
|
|
"max_ttl": "1h",
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("unexpected error: %v", err)
|
|
}
|
|
if resp == nil || !resp.IsError() {
|
|
t.Fatal("expected error when ttl > max_ttl")
|
|
}
|
|
}
|
|
|
|
func TestRole_UnrestrictedModels(t *testing.T) {
|
|
b, s := getTestBackend(t)
|
|
ctx := context.Background()
|
|
|
|
resp, err := b.HandleRequest(ctx, &logical.Request{
|
|
Operation: logical.CreateOperation,
|
|
Path: "roles/open",
|
|
Storage: s,
|
|
Data: map[string]interface{}{"max_budget": 10},
|
|
})
|
|
if err != nil || (resp != nil && resp.IsError()) {
|
|
t.Fatalf("write role: err=%v resp=%v", err, resp)
|
|
}
|
|
role, err := b.getRole(ctx, s, "open")
|
|
if err != nil {
|
|
t.Fatalf("getRole: %v", err)
|
|
}
|
|
if len(role.Models) != 0 {
|
|
t.Fatalf("expected no model restriction, got %v", role.Models)
|
|
}
|
|
}
|