llmgw/internal/router/batches_test.go

package router

import (
	"bytes"
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"mime/multipart"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

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

func upload(t *testing.T, h http.Handler, secret, content string) (int, map[string]any) {
	t.Helper()
	var buf bytes.Buffer
	mw := multipart.NewWriter(&buf)
	_ = mw.WriteField("purpose", "batch")
	fw, _ := mw.CreateFormFile("file", "in.jsonl")
	_, _ = fw.Write([]byte(content))
	_ = mw.Close()
	req := httptest.NewRequest(http.MethodPost, "/v1/files", &buf)
	req.Header.Set("Content-Type", mw.FormDataContentType())
	req.Header.Set("Authorization", "Bearer "+secret)
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	var out map[string]any
	_ = json.Unmarshal(rec.Body.Bytes(), &out)
	return rec.Code, out
}

func call(t *testing.T, h http.Handler, secret, method, path string, body any) (int, map[string]any, string) {
	t.Helper()
	var rd *bytes.Reader
	if body != nil {
		b, _ := json.Marshal(body)
		rd = bytes.NewReader(b)
	} else {
		rd = bytes.NewReader(nil)
	}
	req := httptest.NewRequest(method, path, rd)
	req.Header.Set("Authorization", "Bearer "+secret)
	req.Header.Set("Content-Type", "application/json")
	rec := httptest.NewRecorder()
	h.ServeHTTP(rec, req)
	var out map[string]any
	_ = json.Unmarshal(rec.Body.Bytes(), &out)
	return rec.Code, out, rec.Body.String()
}

func waitBatch(t *testing.T, h http.Handler, secret, id string) map[string]any {
	t.Helper()
	deadline := time.Now().Add(10 * time.Second)
	for {
		_, b, _ := call(t, h, secret, http.MethodGet, "/v1/batches/"+id, nil)
		if store.BatchDone(fmt.Sprint(b["status"])) {
			return b
		}
		if time.Now().After(deadline) {
			t.Fatalf("the batch did not end: %v", b)
		}
		time.Sleep(50 * time.Millisecond)
	}
}

func TestABatchIsServedAndItsAnswersWritten(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := pricedRouter(t, e.srv.URL)
	h := rt.Handler()

	var in strings.Builder
	for i := range 5 {
		fmt.Fprintf(&in, `{"custom_id":"r%d","method":"POST","url":"/v1/chat/completions","body":{"model":"q14","messages":[{"role":"user","content":"hi"}]}}`+"\n", i)
	}
	// One the engine cannot serve: a model that is not there.
	in.WriteString(`{"custom_id":"bad","method":"POST","url":"/v1/chat/completions","body":{"model":"nope","messages":[]}}` + "\n")
	code, f := upload(t, h, secret, in.String())
	if code != 200 || f["purpose"] != "batch" {
		t.Fatalf("upload: %d %v", code, f)
	}
	code, b, raw := call(t, h, secret, http.MethodPost, "/v1/batches", map[string]any{
		"input_file_id": f["id"], "endpoint": "/v1/chat/completions", "completion_window": "24h",
		"metadata": map[string]string{"job": "nightly"}})
	if code != 200 || b["status"] != "validating" {
		t.Fatalf("create: %d %s", code, raw)
	}
	b = waitBatch(t, h, secret, b["id"].(string))
	counts := b["request_counts"].(map[string]any)
	if b["status"] != "completed" || counts["total"] != 6.0 || counts["completed"] != 5.0 || counts["failed"] != 1.0 {
		t.Fatalf("ended %v", b)
	}
	if e.calls.Load() != 5 {
		t.Errorf("the engine served %d", e.calls.Load())
	}
	_, _, out := call(t, h, secret, http.MethodGet, "/v1/files/"+b["output_file_id"].(string)+"/content", nil)
	lines := strings.Split(strings.TrimSpace(out), "\n")
	if len(lines) != 5 || !strings.Contains(lines[0], `"custom_id":"r`) || !strings.Contains(lines[0], `"status_code":200`) {
		t.Errorf("output:\n%s", out)
	}
	_, _, errs := call(t, h, secret, http.MethodGet, "/v1/files/"+b["error_file_id"].(string)+"/content", nil)
	if !strings.Contains(errs, `"custom_id":"bad"`) {
		t.Errorf("errors:\n%s", errs)
	}
	// Charged like any request, and findable by the batch.
	var rid string
	for _, l := range lines {
		var o struct {
			Response struct {
				RequestID string `json:"request_id"`
			} `json:"response"`
		}
		_ = json.Unmarshal([]byte(l), &o)
		rid = o.Response.RequestID
	}
	deadline := time.Now().Add(5 * time.Second)
	for {
		u, ok, _ := rt.store.UsageByRequest(t.Context(), rid)
		if ok {
			if u.Tags["batch"] != b["id"] || u.CostMicros == 0 {
				t.Errorf("usage %+v", u)
			}
			break
		}
		if time.Now().After(deadline) {
			t.Fatal("a batch request left no usage")
		}
		time.Sleep(20 * time.Millisecond)
	}
	_, list, _ := call(t, h, secret, http.MethodGet, "/v1/batches", nil)
	if d := list["data"].([]any); len(d) != 1 {
		t.Errorf("listed %v", list)
	}
}

func TestABatchWithBadLinesFailsAndSaysWhich(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := pricedRouter(t, e.srv.URL)
	h := rt.Handler()
	_, f := upload(t, h, secret, `{"custom_id":"a","method":"POST","url":"/v1/embeddings","body":{"model":"q14"}}
not json
{"custom_id":"a","method":"POST","url":"/v1/chat/completions","body":{"model":"q14"}}
`)
	_, b, _ := call(t, h, secret, http.MethodPost, "/v1/batches", map[string]any{
		"input_file_id": f["id"], "endpoint": "/v1/chat/completions", "completion_window": "24h"})
	b = waitBatch(t, h, secret, b["id"].(string))
	j, _ := json.Marshal(b["errors"])
	if b["status"] != "failed" || !strings.Contains(string(j), "mismatched_endpoint") || !strings.Contains(string(j), "invalid_json_line") {
		t.Errorf("%v %s", b["status"], j)
	}
	if e.calls.Load() != 0 {
		t.Error("a batch that failed validation was served")
	}
}

func TestBatchesAndFilesAreTheirOwners(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := pricedRouter(t, e.srv.URL)
	h := rt.Handler()
	_, f := upload(t, h, secret, `{"custom_id":"a","method":"POST","url":"/v1/chat/completions","body":{"model":"q14"}}`)

	// Another project's key sees none of it.
	org, _ := rt.store.CreateOrg(t.Context(), "other")
	team, _ := rt.store.CreateTeam(t.Context(), org, "t")
	project, _ := rt.store.CreateProject(t.Context(), team, "p")
	other := store.KeyPrefix + strings.Repeat("b", 43)
	sum := sha256.Sum256([]byte(other))
	if err := rt.store.CreateKey(t.Context(), store.Key{ID: "key-2", Hash: sum[:], Prefix: "bbbb", Name: "other",
		ProjectID: project, CreatedAt: time.Now()}); err != nil {
		t.Fatal(err)
	}
	if code, _, _ := call(t, h, other, http.MethodGet, "/v1/files/"+f["id"].(string), nil); code != 404 {
		t.Errorf("another project read a file: %d", code)
	}
	if code, _, _ := call(t, h, other, http.MethodPost, "/v1/batches", map[string]any{
		"input_file_id": f["id"], "endpoint": "/v1/chat/completions", "completion_window": "24h"}); code != 400 {
		t.Errorf("another project batched a file: %d", code)
	}
	if code, _, _ := call(t, h, secret, http.MethodDelete, "/v1/files/"+f["id"].(string), nil); code != 200 {
		t.Errorf("delete: %d", code)
	}
}

func TestABatchPicksUpWhereItStoppedAndCanBeCancelled(t *testing.T) {
	e := newFakeEngine(t, "qwen2.5-14b-instruct")
	rt, secret := pricedRouter(t, e.srv.URL)
	h := rt.Handler()
	var in strings.Builder
	for i := range 4 {
		fmt.Fprintf(&in, `{"custom_id":"r%d","method":"POST","url":"/v1/chat/completions","body":{"model":"q14","messages":[]}}`+"\n", i)
	}
	_, f := upload(t, h, secret, in.String())

	// As a gateway that died after serving the first line left it.
	now := time.Now()
	b := store.Batch{ID: "batch_resumed", Owner: "project:" + mustProjectOf(t, rt, "key-1"), KeyID: "key-1",
		Endpoint: "/v1/chat/completions", InputFileID: f["id"].(string), CompletionWindow: "24h",
		Status: store.BatchInProgress, CreatedAt: now, ExpiresAt: now.Add(24 * time.Hour)}
	if err := rt.store.PutBatch(t.Context(), b); err != nil {
		t.Fatal(err)
	}
	b.Total, b.Completed, b.InProgressAt = 4, 1, now
	_ = rt.store.UpdateBatch(t.Context(), b)
	_ = rt.store.PutBatchResult(t.Context(), b.ID, store.BatchResult{Line: 0, OK: true,
		Data: `{"id":"batch_req_x","custom_id":"r0","response":{"status_code":200},"error":null}`})
	rt.wakeBatches()
	got := waitBatch(t, h, secret, b.ID)
	if c := got["request_counts"].(map[string]any); got["status"] != "completed" || c["completed"] != 4.0 {
		t.Fatalf("%v", got)
	}
	if e.calls.Load() != 3 {
		t.Errorf("served %d; the line kept already was served again", e.calls.Load())
	}

	// Cancelled before it is served: nothing is sent.
	b2 := b
	b2.ID, b2.Status = "batch_cancelled", store.BatchValidating
	if err := rt.store.PutBatch(t.Context(), b2); err != nil {
		t.Fatal(err)
	}
	_, _ = rt.store.CancelBatch(t.Context(), b2.ID, time.Now())
	rt.wakeBatches()
	if got := waitBatch(t, h, secret, b2.ID); got["status"] != "cancelled" {
		t.Errorf("%v", got)
	}
	if e.calls.Load() != 3 {
		t.Errorf("a cancelled batch was served: %d", e.calls.Load())
	}
}

func mustProjectOf(t *testing.T, rt *Router, keyID string) string {
	t.Helper()
	k, err := rt.store.KeyByID(t.Context(), keyID)
	if err != nil {
		t.Fatal(err)
	}
	return k.ProjectID
}