llmgw/internal/router/finish_test.go

package router

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"gitea.mynodes.xyz/ilker/llmgw/internal/accounting"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/admission"
	"gitea.mynodes.xyz/ilker/llmgw/internal/router/proxy"
)

// What a client is told for each way a dispatch can end. Each status is
// something a client acts on: 503 with Retry-After is "try again", 503 alone
// is "not here", 502 is "the engine failed", and a stream already under way
// cannot be given a status at all.
func TestEachWayADispatchEndsIsToldToTheClientAsItShouldBe(t *testing.T) {
	rt := newTestRouter(t, baseConfig())
	for _, c := range []struct {
		name       string
		res        dispatchResult
		status     int
		body       bool // an error body was written
		retryAfter bool
	}{
		{"served", dispatchResult{outcome: proxy.Outcome{Status: 200, Wrote: true}}, 200, false, false},
		{"engine refused, relayed", dispatchResult{outcome: proxy.Outcome{Status: 400, Wrote: true}}, 400, false, false},
		{"client went away", dispatchResult{err: proxy.ErrClientGone}, 499, false, false},
		{"stream cut short", dispatchResult{outcome: proxy.Outcome{Status: 200, Wrote: true}, err: errors.New("eof")}, 200, false, false},
		{"stream cut before a status", dispatchResult{outcome: proxy.Outcome{Wrote: true}, err: errors.New("eof")}, 200, false, false},
		{"still loading", dispatchResult{err: fmt.Errorf("%w: retry shortly", errLoading)}, 503, true, true},
		{"nothing serves it", dispatchResult{err: fmt.Errorf("%w: no replica", errNoCapacity)}, 503, true, false},
		{"engine failed", dispatchResult{err: errors.New("connection reset")}, 502, true, false},
	} {
		t.Run(c.name, func(t *testing.T) {
			w := httptest.NewRecorder()
			got := rt.finish(w, &parsedRequest{Model: "m"}, c.res, time.Now())
			if got != c.status {
				t.Errorf("status %d, expected %d", got, c.status)
			}
			if c.body {
				if w.Code != c.status {
					t.Errorf("written %d, expected %d", w.Code, c.status)
				}
				var e struct {
					Error struct {
						Message string `json:"message"`
					} `json:"error"`
				}
				if err := json.Unmarshal(w.Body.Bytes(), &e); err != nil || e.Error.Message == "" {
					t.Errorf("error body %q", w.Body.String())
				}
			} else if w.Body.Len() != 0 {
				t.Errorf("wrote %q over a response that was not ours to write", w.Body.String())
			}
			if (w.Header().Get("Retry-After") != "") != c.retryAfter {
				t.Errorf("Retry-After %q", w.Header().Get("Retry-After"))
			}
		})
	}
}

func TestAnUpstreamStatusIsKeptWhenTheStreamBreaks(t *testing.T) {
	rt := newTestRouter(t, baseConfig())
	w := httptest.NewRecorder()
	res := dispatchResult{outcome: proxy.Outcome{Status: http.StatusTooManyRequests, Wrote: true}, err: errors.New("eof")}
	if got := rt.finish(w, &parsedRequest{Model: "m"}, res, time.Now()); got != 429 {
		t.Errorf("status %d", got)
	}
}

func TestARefusalBeforeDispatchSaysWhetherWaitingHelps(t *testing.T) {
	rt := newTestRouter(t, baseConfig())
	for _, c := range []struct {
		name       string
		refuse     func(http.ResponseWriter)
		status     int
		retryAfter string
	}{
		{"queue full", func(w http.ResponseWriter) {
			rt.rejectAdmission(w, &parsedRequest{Model: "m"}, admission.ErrQueueFull, time.Now())
		}, 503, ""},
		{"queued too long", func(w http.ResponseWriter) {
			rt.rejectAdmission(w, &parsedRequest{Model: "m"}, admission.ErrQueueTimeout, time.Now())
		}, 504, ""},
		{"displaced for a higher band", func(w http.ResponseWriter) {
			rt.rejectAdmission(w, &parsedRequest{Model: "m"}, admission.ErrDisplaced, time.Now())
		}, 503, "5"},
		{"shutting down", func(w http.ResponseWriter) {
			rt.rejectAdmission(w, &parsedRequest{Model: "m"}, admission.ErrClosed, time.Now())
		}, 503, ""},
		{"client left the queue", func(w http.ResponseWriter) {
			rt.rejectAdmission(w, &parsedRequest{Model: "m"}, context.Canceled, time.Now())
		}, 200, ""}, // nothing written: nobody is listening
		{"rate limited", func(w http.ResponseWriter) {
			rt.rejectLimit(w, &parsedRequest{Model: "m"},
				&accounting.LimitError{Err: accounting.ErrRateLimited, RetryAfter: 3 * time.Second}, time.Now())
		}, 429, "3"},
		{"budget spent", func(w http.ResponseWriter) {
			rt.rejectLimit(w, &parsedRequest{Model: "m"}, &accounting.LimitError{Err: accounting.ErrBudgetSpent}, time.Now())
		}, 429, ""},
		{"not a limit at all", func(w http.ResponseWriter) {
			rt.rejectLimit(w, &parsedRequest{Model: "m"}, errors.New("store down"), time.Now())
		}, 500, ""},
	} {
		t.Run(c.name, func(t *testing.T) {
			w := httptest.NewRecorder()
			c.refuse(w)
			if w.Code != c.status {
				t.Errorf("status %d, expected %d", w.Code, c.status)
			}
			if got := w.Header().Get("Retry-After"); got != c.retryAfter {
				t.Errorf("Retry-After %q, expected %q", got, c.retryAfter)
			}
		})
	}
}