token_thief/tests/logger/queue_test.go

package logger_test

import (
	"context"
	"errors"
	"sync"
	"testing"
	"time"

	"git.misaka.ren/M1saka/token_thief/logger"
)

type batchScript struct {
	prepareErr error
	appendAt   int
	appendErr  error
	sendErr    error
}

type fakeBatch struct {
	script     batchScript
	rows       [][]any
	appendCall int
	abortCalls int
}

func (f *fakeBatch) Append(v ...any) error {
	f.appendCall++
	if f.script.appendErr != nil && f.appendCall == f.script.appendAt {
		return f.script.appendErr
	}
	f.rows = append(f.rows, append([]any(nil), v...))
	return nil
}

func (f *fakeBatch) Send() error  { return f.script.sendErr }
func (f *fakeBatch) Abort() error { f.abortCalls++; return nil }

type fakePool struct {
	calls   int
	queries []string
	batches []*fakeBatch
	scripts []batchScript
}

func (p *fakePool) PrepareBatch(_ context.Context, query string) (logger.Batch, error) {
	idx := p.calls
	p.calls++
	p.queries = append(p.queries, query)
	var script batchScript
	if idx < len(p.scripts) {
		script = p.scripts[idx]
	}
	if script.prepareErr != nil {
		return nil, script.prepareErr
	}
	b := &fakeBatch{script: script}
	p.batches = append(p.batches, b)
	return b, nil
}

func TestFlushSuccessUsesClickHouseTypes(t *testing.T) {
	p := &fakePool{}
	started := time.Date(2026, 7, 9, 1, 2, 3, 4_000_000, time.UTC)
	entry := &logger.LogEntry{
		RequestID: "rid", Method: "POST", Path: "/v1", Query: "a=b", ClientIP: "127.0.0.1",
		RequestHeaders: []byte(`{"x":"y"}`), RequestBody: []byte("request"), RequestTruncated: true,
		StatusCode: 201, ResponseHeaders: []byte(`{"h":"v"}`), ResponseBody: []byte("response"),
		ResponseTruncated: true, IsStream: true, LatencyMS: 150, StartedAt: started,
		FinishedAt: started.Add(150 * time.Millisecond),
	}

	result := logger.Flush(context.Background(), p, []*logger.LogEntry{entry})
	if result.Err != nil || result.Failed != 0 || len(result.Retry) != 0 {
		t.Fatalf("unexpected result: %+v", result)
	}
	if p.calls != 1 || len(p.batches) != 1 || len(p.batches[0].rows) != 1 {
		t.Fatalf("calls=%d batches=%d rows=%d", p.calls, len(p.batches), len(p.batches[0].rows))
	}
	if contains(p.queries[0], "VALUES") {
		t.Fatalf("query contains VALUES: %s", p.queries[0])
	}
	row := p.batches[0].rows[0]
	if len(row) != 17 {
		t.Fatalf("columns=%d want 17", len(row))
	}
	if row[5] != string(entry.RequestHeaders) || row[6] != string(entry.RequestBody) || row[9] != string(entry.ResponseHeaders) || row[10] != string(entry.ResponseBody) {
		t.Fatalf("byte fields were not converted to strings: %#v", row)
	}
	if _, ok := row[8].(int32); !ok {
		t.Fatalf("status code type=%T want int32", row[8])
	}
	if p.batches[0].abortCalls != 0 {
		t.Fatalf("abort calls=%d want 0", p.batches[0].abortCalls)
	}
}

func TestFlushEmptyBatchDoesNotPrepare(t *testing.T) {
	p := &fakePool{}

	result := logger.Flush(context.Background(), p, nil)
	if result.Err != nil || result.Failed != 0 || result.Ambiguous != 0 || len(result.Retry) != 0 {
		t.Fatalf("unexpected result: %+v", result)
	}
	if p.calls != 0 {
		t.Fatalf("PrepareBatch calls=%d want 0", p.calls)
	}
}

func TestFlushPrepareFailureIsRetryable(t *testing.T) {
	p := &fakePool{scripts: []batchScript{{prepareErr: errors.New("prepare")}}}
	entries := []*logger.LogEntry{{RequestID: "a"}, {RequestID: "b"}}
	result := logger.Flush(context.Background(), p, entries)
	if result.Err == nil || result.Failed != 0 || len(result.Retry) != 2 {
		t.Fatalf("unexpected result: %+v", result)
	}
	if len(p.batches) != 0 {
		t.Fatalf("prepare failure created %d batches", len(p.batches))
	}
}

func TestFlushAppendFailureIsolatesOnlyBadRow(t *testing.T) {
	appendErr := errors.New("bad row")
	p := &fakePool{scripts: []batchScript{
		{appendAt: 2, appendErr: appendErr},
		{},
		{appendAt: 1, appendErr: appendErr},
		{},
	}}
	entries := []*logger.LogEntry{{RequestID: "a"}, {RequestID: "bad"}, {RequestID: "c"}}

	result := logger.Flush(context.Background(), p, entries)
	if result.Err == nil || result.Failed != 1 || len(result.Retry) != 0 {
		t.Fatalf("unexpected result: %+v", result)
	}
	if p.calls != 4 {
		t.Fatalf("PrepareBatch calls=%d want 4", p.calls)
	}
	if p.batches[0].abortCalls != 1 || p.batches[2].abortCalls != 1 {
		t.Fatalf("abort calls initial=%d bad-row=%d want 1 each", p.batches[0].abortCalls, p.batches[2].abortCalls)
	}
	if p.batches[1].abortCalls != 0 || p.batches[3].abortCalls != 0 {
		t.Fatalf("successful batches were aborted")
	}
}

func TestFlushAppendFailureCanFullyRecover(t *testing.T) {
	p := &fakePool{scripts: []batchScript{
		{appendAt: 2, appendErr: errors.New("batch append")},
		{},
		{},
	}}
	entries := []*logger.LogEntry{{RequestID: "a"}, {RequestID: "b"}}

	result := logger.Flush(context.Background(), p, entries)
	if result.Err != nil || result.Failed != 0 || len(result.Retry) != 0 {
		t.Fatalf("unexpected result: %+v", result)
	}
	if p.calls != 3 || p.batches[0].abortCalls != 1 {
		t.Fatalf("calls=%d abort=%d", p.calls, p.batches[0].abortCalls)
	}
}

func TestFlushSendFailureIsAmbiguousAndNotRetried(t *testing.T) {
	p := &fakePool{scripts: []batchScript{{sendErr: errors.New("connection lost")}}}
	entries := []*logger.LogEntry{{RequestID: "a"}, {RequestID: "b"}}

	result := logger.Flush(context.Background(), p, entries)
	if result.Err == nil || result.Failed != 2 || len(result.Retry) != 0 {
		t.Fatalf("unexpected result: %+v", result)
	}
	if result.Ambiguous != 2 {
		t.Fatalf("Ambiguous=%d want 2", result.Ambiguous)
	}
	if p.calls != 1 {
		t.Fatalf("PrepareBatch calls=%d want 1", p.calls)
	}
	if p.batches[0].abortCalls != 1 {
		t.Fatalf("abort calls=%d want 1", p.batches[0].abortCalls)
	}
}

func TestFlushAppendRecoveryCountsAmbiguousSingleSend(t *testing.T) {
	p := &fakePool{scripts: []batchScript{
		{appendAt: 1, appendErr: errors.New("bad batch")},
		{},
		{sendErr: errors.New("ack lost")},
	}}
	entries := []*logger.LogEntry{{RequestID: "a"}, {RequestID: "b"}}

	result := logger.Flush(context.Background(), p, entries)
	if result.Failed != 1 || result.Ambiguous != 1 || len(result.Retry) != 0 {
		t.Fatalf("unexpected result: %+v", result)
	}
}

func TestEstimatedBytesCoversStringsAndByteSlices(t *testing.T) {
	entry := &logger.LogEntry{
		RequestID: "1", Method: "22", Path: "333", Query: "4444", ClientIP: "55555", Error: "666666",
		RequestHeaders: []byte("7777777"), RequestBody: []byte("88888888"),
		ResponseHeaders: []byte("999999999"), ResponseBody: []byte("0000000000"),
	}
	if got, want := logger.EstimatedBytes(entry), int64(55); got != want {
		t.Fatalf("EstimatedBytes=%d want %d", got, want)
	}
}

func TestQueueByteBudgetReleasedAfterFinalDiscard(t *testing.T) {
	entry := &logger.LogEntry{
		RequestID: "request", Method: "POST", Path: "/path", Query: "q=1", ClientIP: "ip", Error: "error",
		RequestHeaders: []byte("rh"), RequestBody: []byte("rb"), ResponseHeaders: []byte("sh"), ResponseBody: []byte("sb"),
	}
	budget := logger.EstimatedBytes(entry)
	q := logger.NewQueue(nil, 4, 2, 1, time.Hour, budget)

	q.Submit(entry)
	q.Submit(entry)
	stats := q.Stats()
	if stats.Enqueued != 1 || stats.Dropped != 1 || stats.Bytes != budget {
		t.Fatalf("before drain: %+v budget=%d", stats, budget)
	}
	q.Start(context.Background())
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	if err := q.Stop(ctx); err != nil {
		t.Fatalf("Stop: %v", err)
	}
	stats = q.Stats()
	if stats.Bytes != 0 || stats.Failed != 1 {
		t.Fatalf("after drain: %+v", stats)
	}
}

func TestQueueStopBeforeStartReleasesBudget(t *testing.T) {
	entry := &logger.LogEntry{RequestID: "queued"}
	q := logger.NewQueue(nil, 2, 2, 1, time.Hour, logger.EstimatedBytes(entry))
	q.Submit(entry)

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	if err := q.Stop(ctx); err != nil {
		t.Fatalf("Stop: %v", err)
	}
	if stats := q.Stats(); stats.Bytes != 0 || stats.Failed != 1 {
		t.Fatalf("stats=%+v", stats)
	}
}

func TestQueueCanceledStopEventuallyReleasesAllBudget(t *testing.T) {
	entry := &logger.LogEntry{RequestID: "queued"}
	q := logger.NewQueue(nil, 32, 32, 1, time.Hour, 32*logger.EstimatedBytes(entry))
	q.Start(context.Background())
	for i := 0; i < 32; i++ {
		q.Submit(entry)
	}

	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	_ = q.Stop(ctx)

	deadline := time.Now().Add(time.Second)
	for q.Stats().Bytes != 0 && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if stats := q.Stats(); stats.Bytes != 0 || stats.Failed != 32 {
		t.Fatalf("stats=%+v", stats)
	}
}

func TestQueueStopAndSubmitAreConcurrentAndRepeatSafe(t *testing.T) {
	q := logger.NewQueue(nil, 32, 8, 2, time.Millisecond, 1<<20)
	q.Start(context.Background())

	var submitters sync.WaitGroup
	for i := 0; i < 8; i++ {
		submitters.Add(1)
		go func() {
			defer submitters.Done()
			for j := 0; j < 2_000; j++ {
				q.Submit(&logger.LogEntry{RequestID: "x"})
			}
		}()
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	if err := q.Stop(ctx); err != nil {
		t.Fatalf("first Stop: %v", err)
	}
	submitters.Wait()
	if err := q.Stop(ctx); err != nil {
		t.Fatalf("second Stop: %v", err)
	}
	q.Submit(&logger.LogEntry{RequestID: "after-stop"})
	if stats := q.Stats(); stats.Bytes != 0 {
		t.Fatalf("bytes after Stop=%d want 0", stats.Bytes)
	}
}

func TestStopIsIndependentFromStartContext(t *testing.T) {
	root, cancelRoot := context.WithCancel(context.Background())
	q := logger.NewQueue(nil, 4, 4, 1, time.Hour, 1024)
	q.Start(root)
	q.Submit(&logger.LogEntry{RequestID: "x"})
	cancelRoot()

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	if err := q.Stop(ctx); err != nil {
		t.Fatalf("Stop after root cancellation: %v", err)
	}
	if stats := q.Stats(); stats.Failed != 1 || stats.Bytes != 0 {
		t.Fatalf("stats=%+v", stats)
	}
}

func contains(s, substr string) bool {
	for i := 0; i+len(substr) <= len(s); i++ {
		if s[i:i+len(substr)] == substr {
			return true
		}
	}
	return false
}