localmodelroutex/backend/internal/influx/influx_test.go

package influx

import (
	"context"
	"io"
	"net/http"
	"net/http/httptest"
	"strings"
	"sync"
	"sync/atomic"
	"testing"
	"time"

	"github.com/google/uuid"

	"git.govekar.net/LovroG05/localmodelroutex/internal/store"
)

func TestEncodeEscapesAndSorts(t *testing.T) {
	p := Point{
		Measurement: "my measure,ment",
		Tags:        map[string]string{"z": "last", "a": "has space", "empty": ""},
		Fields: map[string]any{
			"s":   `quote " back \ newline` + "\nend",
			"i":   int(7),
			"i64": int64(-3),
			"f":   1.5,
			"b":   true,
			"nan": 0.0 / zero(),
			"bad": struct{}{},
		},
		Time: time.Unix(1, 2),
	}
	var b strings.Builder
	if !p.encode(&b) {
		t.Fatal("encode returned false")
	}
	want := `my\ measure\,ment,a=has\ space,z=last b=true,f=1.5,i=7i,i64=-3i,s="quote \" back \\ newline end" 1000000002` + "\n"
	if b.String() != want {
		t.Fatalf("got\n%s\nwant\n%s", b.String(), want)
	}
}

func zero() float64 { return 0 }

func TestEncodeNoFields(t *testing.T) {
	var b strings.Builder
	if (Point{Measurement: "m", Fields: map[string]any{"x": struct{}{}}}).encode(&b) {
		t.Fatal("expected false for a point without usable fields")
	}
	if b.Len() != 0 {
		t.Fatalf("expected no output, got %q", b.String())
	}
}

func TestNewRejectsBadConfig(t *testing.T) {
	cases := []Config{
		{Database: "d", Measurement: "m"},
		{URL: "http://x", Measurement: "m"},
		{URL: "http://x", Database: "d"},
		{URL: "not a url", Database: "d", Measurement: "m"},
		{URL: "ftp://x", Database: "d", Measurement: "m"},
	}
	for i, c := range cases {
		if _, err := New(c); err == nil {
			t.Errorf("case %d: expected error for %+v", i, c)
		}
	}
}

func TestNewWriteURL(t *testing.T) {
	e, err := New(Config{URL: "http://influx:8181/prefix/", Database: "lmr", Measurement: "m"})
	if err != nil {
		t.Fatal(err)
	}
	want := "http://influx:8181/prefix/api/v3/write_lp?db=lmr&precision=nanosecond"
	if e.writeURL != want {
		t.Fatalf("writeURL = %q, want %q", e.writeURL, want)
	}
}

// captureServer records every write body and the auth header it saw.
type captureServer struct {
	*httptest.Server
	mu     sync.Mutex
	bodies []string
	auth   string
	status atomic.Int32
	got    chan struct{}
}

func newCaptureServer(t *testing.T) *captureServer {
	t.Helper()
	cs := &captureServer{got: make(chan struct{}, 64)}
	cs.status.Store(http.StatusNoContent)
	cs.Server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if r.URL.Path != "/api/v3/write_lp" {
			t.Errorf("unexpected path %s", r.URL.Path)
		}
		b, _ := io.ReadAll(r.Body)
		cs.mu.Lock()
		cs.bodies = append(cs.bodies, string(b))
		cs.auth = r.Header.Get("Authorization")
		cs.mu.Unlock()
		w.WriteHeader(int(cs.status.Load()))
		cs.got <- struct{}{}
	}))
	t.Cleanup(cs.Close)
	return cs
}

func (cs *captureServer) wait(t *testing.T) {
	t.Helper()
	select {
	case <-cs.got:
	case <-time.After(5 * time.Second):
		t.Fatal("no write arrived")
	}
}

func newTestExporter(t *testing.T, srv *captureServer, cfg Config) *Exporter {
	t.Helper()
	cfg.URL = srv.URL
	cfg.Database = "lmr"
	cfg.Measurement = "lmr_request"
	cfg.Token = "secret"
	e, err := New(cfg)
	if err != nil {
		t.Fatal(err)
	}
	e.now = func() time.Time { return time.Unix(100, 0) }
	return e
}

func TestRecordWritesRequestPoint(t *testing.T) {
	srv := newCaptureServer(t)
	e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 1})
	e.Start()

	uid := uuid.MustParse("11111111-1111-1111-1111-111111111111")
	ttft := 42
	e.Record(store.RequestLog{
		UserID:           &uid,
		KeyPrefix:        "lmr_abc",
		ModelName:        "openai/gpt-4o",
		ProviderName:     "openai",
		StatusCode:       200,
		Streamed:         true,
		PromptTokens:     10,
		CompletionTokens: 5,
		CachedTokens:     2,
		TotalTokens:      15,
		Cost:             0.00125,
		LatencyMS:        321,
		TTFTMs:           &ttft,
		RequestID:        "req-1",
	})
	srv.wait(t)
	if err := e.Close(context.Background()); err != nil {
		t.Fatal(err)
	}

	srv.mu.Lock()
	defer srv.mu.Unlock()
	if srv.auth != "Bearer secret" {
		t.Fatalf("auth = %q", srv.auth)
	}
	want := `lmr_request,fallback=false,key_prefix=lmr_abc,model=openai/gpt-4o,provider=openai,status_code=200,streamed=true,user_id=11111111-1111-1111-1111-111111111111 ` +
		`cached_tokens=2i,completion_tokens=5i,cost=0.00125,error="",errors=0i,finish_reason="",latency_ms=321i,prompt_tokens=10i,request_id="req-1",requests=1i,total_tokens=15i,ttft_ms=42i 100000000000` + "\n"
	if len(srv.bodies) != 1 || srv.bodies[0] != want {
		t.Fatalf("bodies = %q\nwant %q", srv.bodies, want)
	}
}

func TestRecordMarksErrors(t *testing.T) {
	srv := newCaptureServer(t)
	e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 1})
	long := strings.Repeat("x", maxErrorLen+10)
	p := e.requestPoint(store.RequestLog{StatusCode: 502, Error: long})
	if p.Fields["errors"] != 1 {
		t.Fatalf("errors = %v", p.Fields["errors"])
	}
	if got := p.Fields["error"].(string); len(got) != maxErrorLen {
		t.Fatalf("error length = %d", len(got))
	}
	p = e.requestPoint(store.RequestLog{StatusCode: 200, Error: "stream interrupted"})
	if p.Fields["errors"] != 1 {
		t.Fatal("a 200 with an error message must count as an error")
	}
}

// Every field must be present on every point so the columns exist from the
// first write. Optional values get an empty string or zero.
func TestRequestPointWritesAllFields(t *testing.T) {
	srv := newCaptureServer(t)
	e := newTestExporter(t, srv, Config{})
	p := e.requestPoint(store.RequestLog{StatusCode: 200})
	for _, k := range []string{"ttft_ms", "error", "finish_reason", "request_id"} {
		if _, ok := p.Fields[k]; !ok {
			t.Errorf("field %s missing", k)
		}
	}
	if p.Fields["ttft_ms"] != 0 || p.Fields["error"] != "" {
		t.Errorf("fields = %v", p.Fields)
	}
	if p.Tags["key_prefix"] != "(no key)" {
		t.Errorf("key_prefix = %q", p.Tags["key_prefix"])
	}
	if p.Tags["fallback"] != "false" {
		t.Errorf("fallback = %q", p.Tags["fallback"])
	}
	if _, ok := p.Tags["served_model"]; ok {
		t.Errorf("served_model must be absent when the requested model served")
	}
}

func TestRequestPointTagsBackupModel(t *testing.T) {
	srv := newCaptureServer(t)
	e := newTestExporter(t, srv, Config{})
	p := e.requestPoint(store.RequestLog{StatusCode: 200, ModelName: "local/a", ServedModel: "cloud/b"})
	if p.Tags["fallback"] != "true" || p.Tags["served_model"] != "cloud/b" || p.Tags["model"] != "local/a" {
		t.Errorf("tags = %v", p.Tags)
	}
}

func TestBatchAndFlushOnClose(t *testing.T) {
	srv := newCaptureServer(t)
	e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 2})
	e.Start()
	for range 3 {
		e.Record(store.RequestLog{StatusCode: 200})
	}
	srv.wait(t) // first batch of 2
	if err := e.Close(context.Background()); err != nil {
		t.Fatal(err)
	}
	srv.mu.Lock()
	defer srv.mu.Unlock()
	if len(srv.bodies) != 2 {
		t.Fatalf("expected 2 writes, got %d", len(srv.bodies))
	}
	if n := strings.Count(srv.bodies[0], "\n"); n != 2 {
		t.Fatalf("first batch has %d lines, want 2", n)
	}
	if n := strings.Count(srv.bodies[1], "\n"); n != 1 {
		t.Fatalf("second batch has %d lines, want 1", n)
	}
}

func TestFlushOnInterval(t *testing.T) {
	srv := newCaptureServer(t)
	e := newTestExporter(t, srv, Config{FlushInterval: 20 * time.Millisecond, BatchSize: 100})
	e.Start()
	defer e.Close(context.Background()) //nolint:errcheck
	e.Record(store.RequestLog{StatusCode: 200})
	srv.wait(t)
}

func TestEnqueueDropsWhenFull(t *testing.T) {
	srv := newCaptureServer(t)
	e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 1, BufferSize: 1})
	// Not started: nothing drains the queue.
	e.Enqueue(Point{Measurement: "m", Fields: map[string]any{"v": 1}})
	e.Enqueue(Point{Measurement: "m", Fields: map[string]any{"v": 1}})
	if e.dropped.Load() != 1 {
		t.Fatalf("dropped = %d, want 1", e.dropped.Load())
	}
}

func TestWriteDoesNotRetryClientErrors(t *testing.T) {
	srv := newCaptureServer(t)
	srv.status.Store(http.StatusBadRequest)
	e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 1})
	e.write([]Point{{Measurement: "m", Fields: map[string]any{"v": 1}, Time: time.Unix(1, 0)}})
	srv.mu.Lock()
	defer srv.mu.Unlock()
	if len(srv.bodies) != 1 {
		t.Fatalf("expected 1 attempt, got %d", len(srv.bodies))
	}
}

func TestWriteRetriesServerErrors(t *testing.T) {
	srv := newCaptureServer(t)
	srv.status.Store(http.StatusInternalServerError)
	e := newTestExporter(t, srv, Config{FlushInterval: time.Hour, BatchSize: 1})
	e.write([]Point{{Measurement: "m", Fields: map[string]any{"v": 1}, Time: time.Unix(1, 0)}})
	srv.mu.Lock()
	defer srv.mu.Unlock()
	if len(srv.bodies) != writeAttempts {
		t.Fatalf("expected %d attempts, got %d", writeAttempts, len(srv.bodies))
	}
}