diff --git a/backend/internal/service/gateway_record_usage_test.go b/backend/internal/service/gateway_record_usage_test.go index 475dea6f..4e7e545a 100644 --- a/backend/internal/service/gateway_record_usage_test.go +++ b/backend/internal/service/gateway_record_usage_test.go @@ -49,6 +49,31 @@ func newGatewayRecordUsageServiceWithBillingRepoForTest(usageRepo UsageLogReposi return svc } +type openAIRecordUsageBestEffortLogRepoStub struct { + UsageLogRepository + + bestEffortErr error + createErr error + bestEffortCalls int + createCalls int + lastLog *UsageLog + lastCtxErr error +} + +func (s *openAIRecordUsageBestEffortLogRepoStub) CreateBestEffort(ctx context.Context, log *UsageLog) error { + s.bestEffortCalls++ + s.lastLog = log + s.lastCtxErr = ctx.Err() + return s.bestEffortErr +} + +func (s *openAIRecordUsageBestEffortLogRepoStub) Create(ctx context.Context, log *UsageLog) (bool, error) { + s.createCalls++ + s.lastLog = log + s.lastCtxErr = ctx.Err() + return false, s.createErr +} + func TestGatewayServiceRecordUsage_BillingUsesDetachedContext(t *testing.T) { usageRepo := &openAIRecordUsageLogRepoStub{inserted: false, err: context.DeadlineExceeded} userRepo := &openAIRecordUsageUserRepoStub{} diff --git a/backend/internal/service/openai_gateway_record_usage_test.go b/backend/internal/service/openai_gateway_record_usage_test.go index 438e9aeb..cd4d58fd 100644 --- a/backend/internal/service/openai_gateway_record_usage_test.go +++ b/backend/internal/service/openai_gateway_record_usage_test.go @@ -29,31 +29,6 @@ func (s *openAIRecordUsageLogRepoStub) Create(ctx context.Context, log *UsageLog return s.inserted, s.err } -type openAIRecordUsageBestEffortLogRepoStub struct { - UsageLogRepository - - bestEffortErr error - createErr error - bestEffortCalls int - createCalls int - lastLog *UsageLog - lastCtxErr error -} - -func (s *openAIRecordUsageBestEffortLogRepoStub) CreateBestEffort(ctx context.Context, log *UsageLog) error { - s.bestEffortCalls++ - s.lastLog = log - s.lastCtxErr = ctx.Err() - return s.bestEffortErr -} - -func (s *openAIRecordUsageBestEffortLogRepoStub) Create(ctx context.Context, log *UsageLog) (bool, error) { - s.createCalls++ - s.lastLog = log - s.lastCtxErr = ctx.Err() - return false, s.createErr -} - type openAIRecordUsageBillingRepoStub struct { UsageBillingRepository