llmgw/internal/store/capacity_test.go

package store

import (
	"context"
	"path/filepath"
	"testing"
	"time"
)

func capacityStore(t *testing.T) *DB {
	t.Helper()
	s, err := OpenSQLite(context.Background(), filepath.Join(t.TempDir(), "cap.db"))
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { _ = s.Close() })
	return s
}

func TestAMeasurementSurvivesTheTerminalItWasPrintedIn(t *testing.T) {
	s := capacityStore(t)
	at := time.Now().UTC().Truncate(time.Second)

	if err := s.PutCapacity(t.Context(), StoredCapacity{
		Model: "qwen", Profile: "cpu",
		Measurement: []byte(`{"max_concurrency_at_slo":8}`),
		MeasuredAt:  at, MeasuredBy: "ilker",
	}); err != nil {
		t.Fatal(err)
	}

	got, err := s.Capacities(t.Context())
	if err != nil {
		t.Fatal(err)
	}
	if len(got) != 1 {
		t.Fatalf("expected the one measurement, got %d", len(got))
	}
	if got[0].Model != "qwen" || got[0].Profile != "cpu" {
		t.Errorf("wrong row: %+v", got[0])
	}
	if string(got[0].Measurement) != `{"max_concurrency_at_slo":8}` {
		t.Errorf("the measurement came back changed: %s", got[0].Measurement)
	}
	if !got[0].MeasuredAt.UTC().Equal(at) {
		t.Errorf("when it was taken is part of the answer; got %v want %v", got[0].MeasuredAt, at)
	}
}

func TestTheSameModelOnTwoProfilesIsTwoMeasurements(t *testing.T) {
	s := capacityStore(t)
	for _, p := range []string{"cpu", "gpu"} {
		if err := s.PutCapacity(t.Context(), StoredCapacity{
			Model: "qwen", Profile: p, Measurement: []byte(`{}`),
		}); err != nil {
			t.Fatal(err)
		}
	}

	// Capacity belongs to the hardware and the engine arguments as much as to
	// the weights. Keyed by model alone, benchmarking the GPU nodes would
	// overwrite what is known about the CPU ones.
	got, _ := s.Capacities(t.Context())
	if len(got) != 2 {
		t.Fatalf("expected one measurement per profile, got %d", len(got))
	}
}

func TestRemeasuringReplacesRatherThanAccumulates(t *testing.T) {
	s := capacityStore(t)
	older := time.Now().Add(-24 * time.Hour)
	if err := s.PutCapacity(t.Context(), StoredCapacity{
		Model: "qwen", Profile: "cpu",
		Measurement: []byte(`{"max_concurrency_at_slo":4}`), MeasuredAt: older,
	}); err != nil {
		t.Fatal(err)
	}
	if err := s.PutCapacity(t.Context(), StoredCapacity{
		Model: "qwen", Profile: "cpu",
		Measurement: []byte(`{"max_concurrency_at_slo":9}`), MeasuredAt: time.Now(),
	}); err != nil {
		t.Fatal(err)
	}

	// A second run is a better answer to the same question, not another data
	// point: keeping both would leave a reader choosing between them with no
	// basis for the choice.
	got, _ := s.Capacities(t.Context())
	if len(got) != 1 {
		t.Fatalf("expected the measurement to be replaced, got %d rows", len(got))
	}
	if string(got[0].Measurement) != `{"max_concurrency_at_slo":9}` {
		t.Errorf("the older measurement won: %s", got[0].Measurement)
	}
}

func TestMeasurementsComeBackNewestFirst(t *testing.T) {
	s := capacityStore(t)
	now := time.Now()
	_ = s.PutCapacity(t.Context(), StoredCapacity{
		Model: "old", Profile: "cpu", Measurement: []byte(`{}`),
		MeasuredAt: now.Add(-48 * time.Hour),
	})
	_ = s.PutCapacity(t.Context(), StoredCapacity{
		Model: "new", Profile: "cpu", Measurement: []byte(`{}`), MeasuredAt: now,
	})

	got, _ := s.Capacities(t.Context())
	// The old one is the one most likely to describe hardware that has since
	// changed, so it is the one that should need scrolling to.
	if len(got) != 2 || got[0].Model != "new" {
		t.Errorf("expected the newest first, got %+v", got)
	}
}

func TestRubbishIsRefusedBeforeItIsStored(t *testing.T) {
	s := capacityStore(t)
	if err := s.PutCapacity(t.Context(), StoredCapacity{
		Model: "qwen", Profile: "cpu", Measurement: []byte(`not json`),
	}); err == nil {
		t.Error("a measurement that is not JSON should be refused at the write")
	}
	if err := s.PutCapacity(t.Context(), StoredCapacity{
		Profile: "cpu", Measurement: []byte(`{}`),
	}); err == nil {
		t.Error("a measurement with no model should be refused")
	}
}

func TestForgettingAMeasurementThatIsNotThereIsFine(t *testing.T) {
	s := capacityStore(t)
	// The caller wanted it gone, and it is. Erroring would make cleaning up
	// after a removed model a thing that needs a check first.
	if err := s.DeleteCapacity(t.Context(), "ghost", "cpu"); err != nil {
		t.Errorf("deleting an absent measurement should not error: %v", err)
	}
}