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) } }