llmgw/internal/control/adminapi/quotas_test.go

package adminapi

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

	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// fakeQuotaStore is Acme with teams ml and web, and Scania with another ml.
type fakeQuotaStore struct {
	quotas []store.CapacityQuota
}

func (f *fakeQuotaStore) CapacityQuotas(context.Context) ([]store.CapacityQuota, error) {
	return f.quotas, nil
}
func (f *fakeQuotaStore) PutCapacityQuota(_ context.Context, q store.CapacityQuota) error {
	for i := range f.quotas {
		if f.quotas[i].Scope == q.Scope && f.quotas[i].ScopeID == q.ScopeID {
			f.quotas[i] = q
			return nil
		}
	}
	f.quotas = append(f.quotas, q)
	return nil
}
func (f *fakeQuotaStore) DeleteCapacityQuota(_ context.Context, s store.Scope, id string) error {
	for i, q := range f.quotas {
		if q.Scope == s && q.ScopeID == id {
			f.quotas = append(f.quotas[:i], f.quotas[i+1:]...)
		}
	}
	return nil
}
func (f *fakeQuotaStore) ScopeParents(context.Context) (map[string][]string, error) {
	return map[string][]string{"team:t-ml": {"org:o-acme"}, "team:t-web": {"org:o-acme"},
		"team:t-sml": {"org:o-scania"}}, nil
}
func (f *fakeQuotaStore) AllOrgs(context.Context) ([]store.Org, error) {
	return []store.Org{{ID: "o-acme", Name: "Acme"}, {ID: "o-scania", Name: "Scania"}}, nil
}
func (f *fakeQuotaStore) AllTeams(context.Context) ([]store.TeamPath, error) {
	return []store.TeamPath{{ID: "t-ml", Name: "ml", OrgID: "o-acme", OrgName: "Acme"},
		{ID: "t-web", Name: "web", OrgID: "o-acme", OrgName: "Acme"},
		{ID: "t-sml", Name: "ml", OrgID: "o-scania", OrgName: "Scania"}}, nil
}
func (f *fakeQuotaStore) AllProjects(context.Context) ([]store.ProjectPath, error) { return nil, nil }

func TestAScopeIsNamedByPathOrIDAndNeverGuessed(t *testing.T) {
	a, _ := setup(t, testModels())
	a.SetQuotas(&fakeQuotaStore{})
	for in, want := range map[string]string{
		"org:Acme": "o-acme", "org:o-acme": "o-acme", "team:Acme/ml": "t-ml",
		"team: acme / ML ": "t-ml", "team:Scania/ml": "t-sml", "team:web": "t-web",
	} {
		_, id, _, err := a.resolveScope(t.Context(), in)
		if err != nil || id != want {
			t.Errorf("%q: %q %v, expected %q", in, id, err, want)
		}
	}
	for _, bad := range []string{"team:ml", "team:Acme/nope", "group:Acme", "Acme", "org:"} {
		if _, _, _, err := a.resolveScope(t.Context(), bad); err == nil {
			t.Errorf("%q was resolved", bad)
		}
	}
	if _, _, _, err := a.resolveScope(t.Context(), "team:ml"); err == nil || !strings.Contains(err.Error(), "could be any of 2") {
		t.Errorf("an ambiguous name should say so: %v", err)
	}
}

func TestAQuotaIsSetListedAndCleared(t *testing.T) {
	a, _ := setup(t, testModels())
	qs := &fakeQuotaStore{}
	a.SetQuotas(qs)
	mux := http.NewServeMux()
	a.Register(mux)
	call := func(method, path, body string) *httptest.ResponseRecorder {
		req := httptest.NewRequest(method, path, strings.NewReader(body))
		req.Header.Set("Content-Type", "application/json")
		req.Header.Set("X-Llmgw-Actor", "ops")
		rec := httptest.NewRecorder()
		mux.ServeHTTP(rec, req)
		return rec
	}
	if rec := call("PUT", "/api/v1/quotas", `{"scope":"team:Acme/ml","guarantee":{"gpus":8},"limit":{"gpus":16}}`); rec.Code != 200 {
		t.Fatalf("set: %d %s", rec.Code, rec.Body)
	}
	if len(qs.quotas) != 1 || qs.quotas[0].ScopeID != "t-ml" || *qs.quotas[0].Guarantee.GPUs != 8 || qs.quotas[0].UpdatedBy != "ops" {
		t.Errorf("stored %+v", qs.quotas)
	}
	rec := call("GET", "/api/v1/quotas", "")
	var views []QuotaView
	if err := json.Unmarshal(rec.Body.Bytes(), &views); err != nil || len(views) != 1 || views[0].Name != "Acme / ml" {
		t.Errorf("list %d %s", rec.Code, rec.Body)
	}
	if rec := call("PUT", "/api/v1/quotas", `{"scope":"team:Acme/ml"}`); rec.Code != 400 {
		t.Errorf("an empty quota was accepted: %d", rec.Code)
	}
	if rec := call("DELETE", "/api/v1/quotas/team:Acme%2Fml", ""); rec.Code != 200 || len(qs.quotas) != 0 {
		t.Errorf("clear: %d %s %+v", rec.Code, rec.Body, qs.quotas)
	}
}