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) } }