ai-router/internal/server/log_writer_test.go

package server

import (
	"ai-router/internal/protocol"
	"bytes"
	"context"
	"database/sql"
	"log"
	"net/http"
	"net/http/httptest"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"
)

func waitLogCondition(t *testing.T, condition func() bool) {
	t.Helper()
	deadline := time.Now().Add(3 * time.Second)
	for !condition() {
		if time.Now().After(deadline) {
			t.Fatal("condition timed out")
		}
		time.Sleep(time.Millisecond)
	}
}

func logCount(t *testing.T, db *sql.DB) int {
	t.Helper()
	var n int
	if err := db.QueryRow("SELECT COUNT(*) FROM logs").Scan(&n); err != nil {
		t.Fatal(err)
	}
	return n
}

func TestLogAsyncAndRawCopy(t *testing.T) {
	s := testServer(t)
	conn, err := s.db.Conn(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	raw := []byte(`{"input_tokens":12}`)
	done := make(chan struct{})
	go func() {
		s.record(requestLog{Usage: protocol.Usage{Raw: raw}})
		close(done)
	}()
	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("record waited for database")
	}
	raw[0] = 'x'
	conn.Close()
	if err := s.flushLogs(); err != nil {
		t.Fatal(err)
	}
	var got string
	if err := s.db.QueryRow("SELECT raw_usage FROM logs").Scan(&got); err != nil {
		t.Fatal(err)
	}
	if got != `{"input_tokens":12}` {
		t.Fatal(got)
	}
}

func TestLogTimerAndBatchDrain(t *testing.T) {
	s := testServer(t)
	s.record(requestLog{})
	waitLogCondition(t, func() bool { return logCount(t, s.db) == 1 })
	for i := 0; i < 257; i++ {
		s.record(requestLog{})
	}
	if err := s.flushLogs(); err != nil {
		t.Fatal(err)
	}
	if n := logCount(t, s.db); n != 258 {
		t.Fatal(n)
	}
}

func TestLogBackpressureAndConcurrentClose(t *testing.T) {
	path := filepath.Join(t.TempDir(), "logs.db")
	s, err := New(path, "password", nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { s.Close() })
	conn, err := s.db.Conn(context.Background())
	if err != nil {
		t.Fatal(err)
	}
	defer conn.Close()
	s.record(requestLog{})
	waitLogCondition(t, func() bool { return s.db.Stats().WaitCount > 0 })
	if cap(s.logQueue) != 1024 {
		t.Fatal(cap(s.logQueue))
	}
	for i := 0; i < 1024; i++ {
		s.record(requestLog{})
	}
	enqueued := make(chan struct{})
	go func() { s.record(requestLog{}); close(enqueued) }()
	select {
	case <-enqueued:
		t.Fatal("full queue did not apply backpressure")
	case <-time.After(20 * time.Millisecond):
	}
	var wg sync.WaitGroup
	for i := 0; i < 8; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			if err := s.Close(); err != nil {
				t.Error(err)
			}
		}()
	}
	conn.Close()
	done := make(chan struct{})
	go func() { wg.Wait(); close(done) }()
	select {
	case <-done:
	case <-time.After(5 * time.Second):
		t.Fatal("Close deadlocked")
	}
	<-enqueued
	s.record(requestLog{})
	if err := s.flushLogs(); err != nil {
		t.Fatal(err)
	}
	db, err := sql.Open("sqlite", path)
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	if n := logCount(t, db); n != 1026 {
		t.Fatal(n)
	}
}

func TestLogPersistenceErrorIsStickyAndRedacted(t *testing.T) {
	s := testServer(t)
	var output bytes.Buffer
	previous := log.Writer()
	log.SetOutput(&output)
	defer log.SetOutput(previous)
	if _, err := s.db.Exec(`CREATE TRIGGER reject_log BEFORE INSERT ON logs WHEN NEW.key_name='secret' BEGIN SELECT RAISE(FAIL,'private database secret'); END`); err != nil {
		t.Fatal(err)
	}
	s.record(requestLog{KeyName: "ok"})
	s.record(requestLog{KeyName: "secret"})
	start := time.Now()
	err := s.flushLogs()
	if err == nil {
		t.Fatal("missing persistence error")
	}
	if time.Since(start) > time.Second {
		t.Fatal("unbounded retries")
	}
	if n := logCount(t, s.db); n != 0 {
		t.Fatalf("partial transaction: %d", n)
	}
	if strings.Contains(output.String(), "secret") || strings.Contains(err.Error(), "secret") {
		t.Fatal("error leaked metadata")
	}
	if n := strings.Count(output.String(), "retrying"); n != 3 {
		t.Fatalf("retries: %d", n)
	}
	if _, err := s.db.Exec("DROP TRIGGER reject_log"); err != nil {
		t.Fatal(err)
	}
	s.record(requestLog{})
	if err := s.flushLogs(); err == nil {
		t.Fatal("error was not sticky")
	}
	if n := logCount(t, s.db); n != 1 {
		t.Fatal(n)
	}
	if err := s.Close(); err == nil {
		t.Fatal("Close lost persistence error")
	}
	if err := s.Close(); err == nil {
		t.Fatal("second Close lost error")
	}
}

func TestCloseWaitsForActiveRequest(t *testing.T) {
	path := filepath.Join(t.TempDir(), "logs.db")
	s, err := New(path, "test-password", nil)
	if err != nil {
		t.Fatal(err)
	}
	t.Cleanup(func() { s.Close() })
	started, release := make(chan struct{}), make(chan struct{})
	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		close(started)
		<-release
		w.Header().Set("Content-Type", "application/json")
		w.Write([]byte(integrationResponse("openai-chat")))
	}))
	defer up.Close()
	defer func() {
		select {
		case <-release:
		default:
			close(release)
		}
	}()
	cookie := loginCookie(t, s)
	ch := integrationChannel(t, s, cookie, "abc", "openai-chat", up.URL)
	key := integrationCreateKey(t, s, cookie, "active", ch)
	requestDone := make(chan struct{})
	go func() { integrationProxy(s, key, "openai-chat", "abc/gpt-5", false); close(requestDone) }()
	<-started
	closed := make(chan error, 1)
	go func() { closed <- s.Close() }()
	waitLogCondition(t, func() bool { s.lifecycleMu.Lock(); defer s.lifecycleMu.Unlock(); return s.closing })
	if w := call(s, "GET", "/v1/models", "", nil); w.Code != 503 {
		t.Fatal(w.Code)
	}
	select {
	case <-closed:
		t.Fatal("Close skipped active request")
	default:
	}
	close(release)
	<-requestDone
	select {
	case err := <-closed:
		if err != nil {
			t.Fatal(err)
		}
	case <-time.After(3 * time.Second):
		t.Fatal("Close did not drain")
	}
	db, err := sql.Open("sqlite", path)
	if err != nil {
		t.Fatal(err)
	}
	defer db.Close()
	if n := logCount(t, db); n != 1 {
		t.Fatal(n)
	}
}