llmgw/internal/control/facility_test.go

package control

import (
	"context"
	"log/slog"
	"maps"
	"slices"
	"strings"
	"testing"

	"gitea.mynodes.xyz/ilker/llmgw/internal/control/facility"
	"gitea.mynodes.xyz/ilker/llmgw/internal/control/inventory"
	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

// memFacilityStore is the store's two tables in memory.
type memFacilityStore struct {
	holds  map[string]store.FacilityHold
	labels map[string]map[string]string
}

func newMemFacilityStore() *memFacilityStore {
	return &memFacilityStore{holds: map[string]store.FacilityHold{}, labels: map[string]map[string]string{}}
}

func (m *memFacilityStore) FacilityHolds(context.Context) ([]store.FacilityHold, error) {
	return slices.Collect(maps.Values(m.holds)), nil
}
func (m *memFacilityStore) SetFacilityHold(_ context.Context, h store.FacilityHold) error {
	m.holds[h.Domain] = h
	return nil
}
func (m *memFacilityStore) DeleteFacilityHold(_ context.Context, d string) error {
	delete(m.holds, d)
	return nil
}
func (m *memFacilityStore) NodeLabels(context.Context) (map[string]map[string]string, error) {
	return m.labels, nil
}
func (m *memFacilityStore) SetNodeLabel(_ context.Context, node, k, v, _ string) error {
	if m.labels[node] == nil {
		m.labels[node] = map[string]string{}
	}
	if v == "" {
		delete(m.labels[node], k)
	} else {
		m.labels[node][k] = v
	}
	return nil
}

func rackNode(id, rack string) inventory.Node {
	return inventory.Node{ID: id, State: inventory.Ready, Labels: map[string]string{"llmgw.io/rack": rack}}
}

func TestADrainingHoldDrainsTheRackAndReleasingItReadiesItAgain(t *testing.T) {
	inv := inventory.NewStore()
	inv.Upsert(rackNode("a1", "r1"))
	inv.Upsert(rackNode("a2", "r1"))
	inv.Upsert(rackNode("b1", "r2"))
	f := NewFacility(inv, slog.New(slog.DiscardHandler))
	ctx := t.Context()

	drained, err := f.Hold(ctx, facility.Hold{Domain: "rack:r1", Reason: "PDU swap", Drain: true})
	if err != nil || !slices.Equal(drained, []string{"a1", "a2"}) {
		t.Fatalf("drained %v, %v", drained, err)
	}
	if n, _ := inv.Get("b1"); n.State != inventory.Ready {
		t.Errorf("b1, in another rack, is %s", n.State)
	}
	if _, held := inv.Facility().HeldIn([]facility.Domain{"rack:r1"}); !held {
		t.Error("the hold is not in force for placement")
	}

	_, readied, err := f.Release(ctx, "rack:r1")
	if err != nil || !slices.Equal(readied, []string{"a1", "a2"}) {
		t.Fatalf("readied %v, %v", readied, err)
	}
	if _, _, err := f.Release(ctx, "rack:r1"); err == nil {
		t.Error("releasing a rack not held is not an error")
	}
}

func TestAHoldWithoutDrainLeavesTheNodesAsTheyAre(t *testing.T) {
	inv := inventory.NewStore()
	inv.Upsert(rackNode("a1", "r1"))
	f := NewFacility(inv, nil)
	if drained, _ := f.Hold(t.Context(), facility.Hold{Domain: "rack:r1"}); len(drained) != 0 {
		t.Errorf("drained %v without being asked", drained)
	}
	if n, _ := inv.Get("a1"); n.State != inventory.Ready {
		t.Errorf("a1 is %s", n.State)
	}
}

func TestHoldsAndLabelsAreKeptAndReadBackByTheNextControlPlane(t *testing.T) {
	st := newMemFacilityStore()
	inv := inventory.NewStore()
	f := NewFacility(inv, nil)
	f.UseStore(t.Context(), st)
	if _, err := f.Hold(t.Context(), facility.Hold{Domain: "pdu:p1", Reason: "breaker", By: "ilker"}); err != nil {
		t.Fatal(err)
	}
	if _, err := f.SetLabels(t.Context(), "gpu01", map[string]string{"llmgw.io/rack": "r9"}, nil, "ilker"); err != nil {
		t.Fatal(err)
	}

	// A restart: a new inventory and facility over the same store.
	inv2 := inventory.NewStore()
	f2 := NewFacility(inv2, nil)
	f2.SetPolicy(facility.Policy{Margin: 0.1})
	f2.UseStore(t.Context(), st)
	inv2.Upsert(inventory.Node{ID: "gpu01", State: inventory.Ready})
	if n, _ := inv2.Get("gpu01"); n.Labels["llmgw.io/rack"] != "r9" {
		t.Errorf("the label set before the restart is gone: %v", n.Labels)
	}
	pol := inv2.Facility()
	if h, ok := pol.Holds["pdu:p1"]; !ok || h.Reason != "breaker" || pol.Margin != 0.1 {
		t.Errorf("after the restart: holds %v, margin %v", pol.Holds, pol.Margin)
	}
}

func TestAMisspeltLocationLabelIsRefused(t *testing.T) {
	f := NewFacility(inventory.NewStore(), nil)
	_, err := f.SetLabels(t.Context(), "n1", map[string]string{"llmgw.io/rak": "r1"}, nil, "")
	if err == nil || !strings.Contains(err.Error(), "llmgw.io/site, room, row, rack or pdu") {
		t.Errorf("got %v", err)
	}
	if _, err := f.SetLabels(t.Context(), "n1", map[string]string{"team": "nlp"}, nil, ""); err != nil {
		t.Errorf("a free label was refused: %v", err)
	}
}