llmgw/internal/store/batches_test.go

package store

import (
	"errors"
	"os"
	"path/filepath"
	"strings"
	"testing"
	"time"

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

func TestFilesAndBatches(t *testing.T) {
	checkFilesAndBatches(t, capacityStore(t))
}

func TestPostgresFilesAndBatches(t *testing.T) {
	checkFilesAndBatches(t, openPG(t))
}

func checkFilesAndBatches(t *testing.T, db *DB) {
	ctx := t.Context()
	now := time.Now()
	if err := db.PutFile(ctx, File{ID: "file-1", Owner: "p1", Purpose: "batch", Filename: "in.jsonl", CreatedAt: now},
		[]byte("{\"a\":1}\n")); err != nil {
		t.Fatal(err)
	}
	if err := db.PutFile(ctx, File{ID: "file-old", Owner: "p1", Purpose: "batch_output", Filename: "out.jsonl",
		CreatedAt: now, ExpiresAt: now.Add(-time.Minute)}, []byte("x")); err != nil {
		t.Fatal(err)
	}
	f, err := db.FileByID(ctx, "file-1")
	if err != nil || f.Bytes != 8 || f.Owner != "p1" {
		t.Fatalf("%+v %v", f, err)
	}
	if b, _ := db.FileContent(ctx, "file-1"); string(b) != "{\"a\":1}\n" {
		t.Errorf("content %q", b)
	}
	if fs, _ := db.Files(ctx, "p1", "batch", 0); len(fs) != 1 {
		t.Errorf("listed %v", fs)
	}
	if n, _ := db.DeleteExpiredFiles(ctx, now); n != 1 {
		t.Errorf("expired %d", n)
	}

	b := Batch{ID: "batch-1", Owner: "p1", KeyID: "k", Endpoint: "/v1/chat/completions", InputFileID: "file-1",
		CompletionWindow: "24h", Status: BatchValidating, CreatedAt: now, ExpiresAt: now.Add(24 * time.Hour)}
	if err := db.PutBatch(ctx, b); err != nil {
		t.Fatal(err)
	}
	got, ok, err := db.ClaimBatch(ctx, "gw-a", now, now.Add(time.Minute))
	if err != nil || !ok || got.ID != "batch-1" {
		t.Fatalf("claim: %+v %v %v", got, ok, err)
	}
	if _, ok, _ := db.ClaimBatch(ctx, "gw-b", now, now.Add(time.Minute)); ok {
		t.Error("a leased batch was claimed twice")
	}
	if held, _ := db.RenewBatchLease(ctx, "batch-1", "gw-b", now.Add(time.Hour)); held {
		t.Error("another gateway renewed the lease")
	}
	got.Status, got.Total, got.InProgressAt = BatchInProgress, 2, now
	if err := db.UpdateBatch(ctx, got); err != nil {
		t.Fatal(err)
	}
	for _, r := range []BatchResult{{Line: 1, OK: false, Data: "e"}, {Line: 0, OK: true, Data: "a"}, {Line: 0, OK: false, Data: "dup"}} {
		if err := db.PutBatchResult(ctx, "batch-1", r); err != nil {
			t.Fatal(err)
		}
	}
	rs, _ := db.BatchResults(ctx, "batch-1")
	if len(rs) != 2 || rs[0].Data != "a" || !rs[0].OK || rs[1].OK {
		t.Errorf("results %+v", rs)
	}
	c, err := db.CancelBatch(ctx, "batch-1", now)
	if err != nil || c.Status != BatchCancelling || c.Total != 2 || c.CancellingAt.IsZero() {
		t.Errorf("cancel %+v %v", c, err)
	}
	if bs, _ := db.Batches(ctx, "p1", "", 0); len(bs) != 1 {
		t.Errorf("listed %d", len(bs))
	}
	if _, err := db.BatchByID(ctx, "nope"); !errors.Is(err, ErrNotFound) {
		t.Errorf("missing: %v", err)
	}
}

func TestClusterSecretsAreSealedUnderAMasterKey(t *testing.T) {
	db := capacityStore(t)
	ctx := t.Context()
	plain, err := db.ClusterSecret(ctx, "old")
	if err != nil {
		t.Fatal(err)
	}
	key := filepath.Join(t.TempDir(), "master")
	_ = os.WriteFile(key, []byte(strings.Repeat("ab", 32)), 0o600)
	sealer, err := secrets.NewSealer(ctx, key)
	if err != nil {
		t.Fatal(err)
	}
	db.SetSealer(sealer)
	if again, err := db.ClusterSecret(ctx, "old"); err != nil || again != plain {
		t.Fatalf("a secret from before changed: %q %v", again, err)
	}
	fresh, err := db.ClusterSecret(ctx, "new")
	if err != nil {
		t.Fatal(err)
	}
	for _, name := range []string{"old", "new"} {
		var v string
		_ = db.db.QueryRowContext(ctx, `SELECT value FROM cluster_secrets WHERE name = ?`, name).Scan(&v)
		if !secrets.IsSealed(v) {
			t.Errorf("%s is kept plain: %q", name, v)
		}
	}
	if again, _ := db.ClusterSecret(ctx, "new"); again != fresh {
		t.Error("a sealed secret read back differently")
	}
	db.SetSealer(nil)
	if _, err := db.ClusterSecret(ctx, "new"); err == nil {
		t.Error("a sealed secret was read without its key")
	}
}