gpt-load/internal/requestlog/test_helpers_test.go

package requestlog

import (
	"context"
	"fmt"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	gormsqlite "github.com/glebarez/sqlite"
	"gorm.io/gorm"
	"gorm.io/gorm/logger"

	"gpt-load/internal/channel"
	"gpt-load/internal/execution"
	"gpt-load/internal/platform/redact"
	"gpt-load/internal/pricing"
	"gpt-load/internal/protocol"
	"gpt-load/internal/storage"
	"gpt-load/internal/storage/models"
	"gpt-load/internal/telemetry"
	"gpt-load/internal/usage"
)

type staticRetentionPolicy struct{ days int }

func (policy staticRetentionPolicy) RequestLogRetentionDays() int {
	return policy.days
}

func newRequestLogTestService(db *gorm.DB) *Service {
	return NewService(
		db,
		redact.New(),
		staticRetentionPolicy{days: 7},
	)
}

type batchWriterFunc func(context.Context, []models.RequestLog) error

func (fn batchWriterFunc) WriteBatch(ctx context.Context, rows []models.RequestLog) error {
	return fn(ctx, rows)
}

type manualTimer struct {
	ch   chan time.Time
	once sync.Once
}

func newManualTimer() *manualTimer {
	return &manualTimer{ch: make(chan time.Time, 1)}
}

func (timer *manualTimer) C() <-chan time.Time {
	return timer.ch
}

func (timer *manualTimer) Stop() bool {
	stopped := false
	timer.once.Do(func() {
		stopped = true
	})
	return stopped
}

func (timer *manualTimer) Fire() {
	timer.ch <- time.Unix(1, 0)
}

type manualTimerFactory struct {
	created chan *manualTimer
}

func newManualTimerFactory() *manualTimerFactory {
	return &manualTimerFactory{created: make(chan *manualTimer, 32)}
}

func (factory *manualTimerFactory) New(time.Duration) workerTimer {
	timer := newManualTimer()
	factory.created <- timer
	return timer
}

func receiveValue[T any](t *testing.T, ch <-chan T) T {
	t.Helper()

	select {
	case value := <-ch:
		return value
	case <-time.After(5 * time.Second):
		t.Fatal("timed out waiting for test coordination")
		var zero T
		return zero
	}
}

func waitGroupDone(t *testing.T, group *sync.WaitGroup) {
	t.Helper()

	done := make(chan struct{})
	go func() {
		group.Wait()
		close(done)
	}()
	receiveValue(t, done)
}

func openRequestLogFileDB(t *testing.T) (*gorm.DB, string) {
	t.Helper()
	dsn := filepath.Join(t.TempDir(), "requestlog.db")
	db, err := storage.Open(dsn)
	if err != nil {
		t.Fatalf("storage.Open(%q) error = %v", dsn, err)
	}
	sqlDB, err := db.DB()
	if err != nil {
		t.Fatalf("db.DB() error = %v", err)
	}
	t.Cleanup(func() {
		if err := sqlDB.Close(); err != nil {
			t.Errorf("close request log database: %v", err)
		}
	})
	if err := storage.AutoMigrate(db); err != nil {
		t.Fatalf("storage.AutoMigrate() error = %v", err)
	}
	return db, dsn
}

func holdRequestLogRollbackJournalReadLock(
	t *testing.T,
	appDB *gorm.DB,
	dsn string,
) func() {
	t.Helper()
	if err := appDB.Exec("PRAGMA busy_timeout = 1").Error; err != nil {
		t.Fatalf("set app busy_timeout: %v", err)
	}
	var mode string
	if err := appDB.Raw("PRAGMA journal_mode = DELETE").Scan(&mode).Error; err != nil {
		t.Fatalf("set rollback journal: %v", err)
	}
	if !strings.EqualFold(mode, "delete") {
		t.Fatalf("journal_mode = %q, want delete", mode)
	}

	blocker, err := gorm.Open(
		gormsqlite.Open(dsn+"?_pragma=busy_timeout(1)"),
		&gorm.Config{Logger: logger.Default.LogMode(logger.Silent)},
	)
	if err != nil {
		t.Fatalf("open blocker: %v", err)
	}
	blockerSQL, err := blocker.DB()
	if err != nil {
		t.Fatalf("blocker DB(): %v", err)
	}
	readTx := blocker.Begin()
	if readTx.Error != nil {
		t.Fatal(readTx.Error)
	}
	var count int64
	if err := readTx.Table("request_logs").Count(&count).Error; err != nil {
		t.Fatal(err)
	}

	var once sync.Once
	release := func() {
		once.Do(func() {
			if err := readTx.Rollback().Error; err != nil {
				t.Errorf("release read lock: %v", err)
			}
			if err := blockerSQL.Close(); err != nil {
				t.Errorf("close blocker: %v", err)
			}
		})
	}
	t.Cleanup(release)
	return release
}

func aggregationRow(
	id string,
	completedAt time.Time,
	groupID uint,
	model string,
) models.RequestLog {
	return models.RequestLog{
		ID:                   id,
		CompletedAtMS:        completedAt.UTC().UnixMilli(),
		AccessKeyID:          1,
		GroupID:              groupID,
		Protocol:             string(protocol.OpenAICompletions),
		ClientModel:          model,
		UpstreamModel:        model,
		Status:               string(telemetry.RequestStatusSuccess),
		StatusCode:           200,
		DurationMs:           10,
		AttemptCount:         1,
		UncachedInputTokens:  1,
		OutputTokens:         2,
		EstimatedCostNanoUSD: 250_000_000,
		UsageState:           string(usage.StateComplete),
		CostState:            string(pricing.CostStatePriced),
		PricingCompleteness:  string(pricing.CompletenessComplete),
	}
}

func aggregationRequestID(index int) string {
	return fmt.Sprintf("00000000-0000-4000-8000-%012d", 900000+index)
}

func testEvent(id string) telemetry.RequestEvent {
	return telemetry.RequestEvent{
		RequestID:             id,
		CompletedAt:           time.Date(2026, time.July, 24, 12, 0, 0, 0, time.UTC),
		AccessKeyID:           42,
		ClientModel:           "client-model",
		UpstreamModel:         "upstream-model",
		UpstreamReportedModel: "upstream-model",
		ModelConsistency:      telemetry.ModelConsistencyMatch,
		Status:                telemetry.RequestStatusSuccess,
		StatusCode:            200,
		DurationMs:            25,
		Attempts: []telemetry.Attempt{{
			Sequence:        1,
			GroupID:         7,
			GroupName:       "primary",
			ChannelID:       channel.OpenAI,
			CredentialID:    8,
			Operation:       execution.OperationChatCompletion,
			RouteMode:       channel.RouteNative,
			UpstreamModel:   "upstream-model",
			DispatchState:   execution.DispatchMaybeSent,
			StatusCode:      200,
			DurationMs:      20,
			FailureCategory: telemetry.FailureCategoryOK,
			Action:          telemetry.ActionTerminate,
		}},
		Usage: telemetry.UsageObservation{
			GroupID:         7,
			ChannelID:       channel.OpenAI,
			CredentialID:    8,
			AttemptSequence: 1,
			Result: usage.Result{
				State: usage.StateNotApplicable,
			},
			Pricing: telemetry.PricingObservation{
				UpstreamModel:       "upstream-model",
				CostState:           string(pricing.CostStateNotApplicable),
				PricingCompleteness: string(pricing.CompletenessNotApplicable),
			},
		},
	}
}

func mustMapEvent(
	t testing.TB,
	redactor *redact.Redactor,
	event telemetry.RequestEvent,
	_ ...*pricing.Table,
) models.RequestLog {
	t.Helper()
	row, err := mapEvent(redactor, event)
	if err != nil {
		t.Fatalf("mapEvent() error = %v", err)
	}
	return row
}