terraform-provider-litellmvaultsecret/internal/provider/client_test.go

package provider

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"testing"
)

// mockVault is a minimal in-memory fake of the Vault/OpenBao HTTP API covering
// the endpoints this provider uses.
type mockVault struct {
	server *httptest.Server

	mu      sync.Mutex
	mounts  map[string]map[string]interface{} // "litellm/" -> mount output
	configs map[string]map[string]interface{} // backend -> config data
	roles   map[string]map[string]interface{} // "backend/roles/name" -> role data
}

func newMockVault(t *testing.T) (*mockVault, *vaultClient) {
	t.Helper()
	m := &mockVault{
		mounts:  map[string]map[string]interface{}{},
		configs: map[string]map[string]interface{}{},
		roles:   map[string]map[string]interface{}{},
	}
	m.server = httptest.NewServer(http.HandlerFunc(m.handle))
	t.Cleanup(m.server.Close)

	client, err := newVaultClient(m.server.URL, "test-token")
	if err != nil {
		t.Fatalf("newVaultClient: %v", err)
	}
	return m, client
}

func (m *mockVault) handle(w http.ResponseWriter, r *http.Request) {
	m.mu.Lock()
	defer m.mu.Unlock()

	path := strings.TrimPrefix(r.URL.Path, "/v1/")

	switch {
	case path == "sys/mounts" && r.Method == http.MethodGet:
		writeSecret(w, map[string]interface{}{"data": toIface(m.mounts)})

	case strings.HasPrefix(path, "sys/mounts/"):
		mp := strings.TrimPrefix(path, "sys/mounts/")
		key := strings.TrimRight(mp, "/") + "/"
		switch r.Method {
		case http.MethodPost, http.MethodPut:
			var body map[string]interface{}
			_ = json.NewDecoder(r.Body).Decode(&body)
			// Store the mount *output* shape (Vault returns int TTLs, not the
			// string TTLs it accepts on input).
			m.mounts[key] = map[string]interface{}{
				"type":        body["type"],
				"description": body["description"],
				"config":      map[string]interface{}{},
			}
			w.WriteHeader(http.StatusNoContent)
		case http.MethodDelete:
			delete(m.mounts, key)
			w.WriteHeader(http.StatusNoContent)
		}

	case strings.HasSuffix(path, "/config"):
		backend := strings.TrimSuffix(path, "/config")
		switch r.Method {
		case http.MethodPost, http.MethodPut:
			var body map[string]interface{}
			_ = json.NewDecoder(r.Body).Decode(&body)
			m.configs[backend] = body
			w.WriteHeader(http.StatusNoContent)
		case http.MethodGet:
			cfg, ok := m.configs[backend]
			if !ok {
				w.WriteHeader(http.StatusNotFound)
				return
			}
			// The real backend never returns master_key.
			out := map[string]interface{}{}
			for k, v := range cfg {
				if k == "master_key" {
					continue
				}
				out[k] = v
			}
			writeSecret(w, map[string]interface{}{"data": out})
		}

	case strings.Contains(path, "/roles/"):
		switch r.Method {
		case http.MethodPost, http.MethodPut:
			var body map[string]interface{}
			_ = json.NewDecoder(r.Body).Decode(&body)
			m.roles[path] = body
			w.WriteHeader(http.StatusNoContent)
		case http.MethodGet:
			role, ok := m.roles[path]
			if !ok {
				w.WriteHeader(http.StatusNotFound)
				return
			}
			writeSecret(w, map[string]interface{}{"data": role})
		case http.MethodDelete:
			delete(m.roles, path)
			w.WriteHeader(http.StatusNoContent)
		}

	default:
		w.WriteHeader(http.StatusNotFound)
	}
}

func toIface(m map[string]map[string]interface{}) map[string]interface{} {
	out := map[string]interface{}{}
	for k, v := range m {
		out[k] = v
	}
	return out
}

func writeSecret(w http.ResponseWriter, body map[string]interface{}) {
	w.Header().Set("Content-Type", "application/json")
	w.WriteHeader(http.StatusOK)
	_ = json.NewEncoder(w).Encode(body)
}

func TestClient_MountLifecycle(t *testing.T) {
	_, c := newMockVault(t)
	ctx := context.Background()

	if err := c.enableMount(ctx, "litellm", "vault-plugin-secrets-litellm", "desc", mountConfig{}); err != nil {
		t.Fatalf("enableMount: %v", err)
	}
	info, err := c.mountInfo(ctx, "litellm")
	if err != nil {
		t.Fatalf("mountInfo: %v", err)
	}
	if info == nil {
		t.Fatal("expected mount to exist")
	}
	if info.Type != "vault-plugin-secrets-litellm" {
		t.Fatalf("unexpected mount type: %q", info.Type)
	}

	if err := c.disableMount(ctx, "litellm"); err != nil {
		t.Fatalf("disableMount: %v", err)
	}
	info, err = c.mountInfo(ctx, "litellm")
	if err != nil {
		t.Fatalf("mountInfo after delete: %v", err)
	}
	if info != nil {
		t.Fatal("expected mount to be gone")
	}
}

func TestClient_ConfigRoundTrip(t *testing.T) {
	m, c := newMockVault(t)
	ctx := context.Background()

	err := c.writeConfig(ctx, "litellm", map[string]interface{}{
		"base_url":                "http://litellm:4000",
		"master_key":              "sk-secret",
		"request_timeout_seconds": 30,
	})
	if err != nil {
		t.Fatalf("writeConfig: %v", err)
	}

	cfg, err := c.readConfig(ctx, "litellm")
	if err != nil {
		t.Fatalf("readConfig: %v", err)
	}
	if cfg["base_url"] != "http://litellm:4000" {
		t.Fatalf("unexpected base_url: %v", cfg["base_url"])
	}
	if _, leaked := cfg["master_key"]; leaked {
		t.Fatal("master_key must not be returned by read")
	}
	_ = m
}

func TestClient_ReadConfigMissing(t *testing.T) {
	_, c := newMockVault(t)
	cfg, err := c.readConfig(context.Background(), "litellm")
	if err != nil {
		t.Fatalf("readConfig: %v", err)
	}
	if cfg != nil {
		t.Fatalf("expected nil config, got %v", cfg)
	}
}

func TestClient_RoleLifecycle(t *testing.T) {
	_, c := newMockVault(t)
	ctx := context.Background()

	err := c.writeRole(ctx, "litellm", "team-a", map[string]interface{}{
		"models":     []string{"gpt-4"},
		"max_budget": 25.0,
		"ttl":        3600,
	})
	if err != nil {
		t.Fatalf("writeRole: %v", err)
	}

	role, err := c.readRole(ctx, "litellm", "team-a")
	if err != nil {
		t.Fatalf("readRole: %v", err)
	}
	if role == nil {
		t.Fatal("expected role to exist")
	}

	if err := c.deleteRole(ctx, "litellm", "team-a"); err != nil {
		t.Fatalf("deleteRole: %v", err)
	}
	role, err = c.readRole(ctx, "litellm", "team-a")
	if err != nil {
		t.Fatalf("readRole after delete: %v", err)
	}
	if role != nil {
		t.Fatal("expected role to be gone")
	}
}

func TestRolePath(t *testing.T) {
	if got := rolePath("litellm/", "team-a"); got != "litellm/roles/team-a" {
		t.Fatalf("unexpected role path: %q", got)
	}
}