You've already forked agentic-coding-workflow
99 lines
3.2 KiB
Go
99 lines
3.2 KiB
Go
package chat
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
// usageRepo is a tiny stub that records what was called and returns canned values.
|
|
// It satisfies only the UsageService-relevant subset of Repository; unused methods
|
|
// inherit panic stubs from fakeRepo so we don't need to declare them here.
|
|
type usageRepo struct {
|
|
*fakeRepo
|
|
|
|
insertedUsage []UsageRecord
|
|
|
|
byUserCalls []timeRange
|
|
byModelCalls []timeRange
|
|
byDayCalls []timeRange
|
|
dailyCalls []userTimeRange
|
|
|
|
usersOut []UserUsageRow
|
|
modelsOut []ModelUsageRow
|
|
daysOut []DayUsageRow
|
|
}
|
|
|
|
type timeRange struct{ from, to time.Time }
|
|
type userTimeRange struct {
|
|
user uuid.UUID
|
|
from time.Time
|
|
to time.Time
|
|
}
|
|
|
|
func (r *usageRepo) InsertUsage(_ context.Context, rec UsageRecord) error {
|
|
r.insertedUsage = append(r.insertedUsage, rec)
|
|
return nil
|
|
}
|
|
func (r *usageRepo) SummarizeUsageByUser(_ context.Context, from, to time.Time) ([]UserUsageRow, error) {
|
|
r.byUserCalls = append(r.byUserCalls, timeRange{from, to})
|
|
return r.usersOut, nil
|
|
}
|
|
func (r *usageRepo) SummarizeUsageByModel(_ context.Context, from, to time.Time) ([]ModelUsageRow, error) {
|
|
r.byModelCalls = append(r.byModelCalls, timeRange{from, to})
|
|
return r.modelsOut, nil
|
|
}
|
|
func (r *usageRepo) SummarizeUsageByDay(_ context.Context, from, to time.Time) ([]DayUsageRow, error) {
|
|
r.byDayCalls = append(r.byDayCalls, timeRange{from, to})
|
|
return r.daysOut, nil
|
|
}
|
|
func (r *usageRepo) SummarizeUserDaily(_ context.Context, uid uuid.UUID, from, to time.Time) ([]DayUsageRow, error) {
|
|
r.dailyCalls = append(r.dailyCalls, userTimeRange{uid, from, to})
|
|
return r.daysOut, nil
|
|
}
|
|
|
|
func TestUsageService_DelegatesToRepo(t *testing.T) {
|
|
repo := &usageRepo{fakeRepo: newFakeRepo()}
|
|
svc := NewUsageService(repo)
|
|
from := time.Date(2026, 5, 1, 0, 0, 0, 0, time.UTC)
|
|
to := from.Add(24 * time.Hour)
|
|
uid := uuid.Must(uuid.NewV7())
|
|
|
|
rec := UsageRecord{
|
|
ConversationID: uuid.Must(uuid.NewV7()),
|
|
MessageID: 42,
|
|
UserID: uid,
|
|
PromptTokens: 100, CompletionTokens: 50, ThinkingTokens: 0,
|
|
CostUSD: 0.001,
|
|
}
|
|
require.NoError(t, svc.Insert(context.Background(), rec))
|
|
require.Len(t, repo.insertedUsage, 1)
|
|
require.Equal(t, rec, repo.insertedUsage[0])
|
|
|
|
repo.usersOut = []UserUsageRow{{UserID: uid, PromptTokens: 100}}
|
|
users, err := svc.SummarizeByUser(context.Background(), from, to)
|
|
require.NoError(t, err)
|
|
require.Equal(t, repo.usersOut, users)
|
|
require.Equal(t, []timeRange{{from, to}}, repo.byUserCalls)
|
|
|
|
repo.modelsOut = []ModelUsageRow{{ModelID: uuid.Must(uuid.NewV7()), CostUSD: 1.5}}
|
|
models, err := svc.SummarizeByModel(context.Background(), from, to)
|
|
require.NoError(t, err)
|
|
require.Equal(t, repo.modelsOut, models)
|
|
require.Equal(t, []timeRange{{from, to}}, repo.byModelCalls)
|
|
|
|
repo.daysOut = []DayUsageRow{{Day: from, PromptTokens: 200}}
|
|
days, err := svc.SummarizeByDay(context.Background(), from, to)
|
|
require.NoError(t, err)
|
|
require.Equal(t, repo.daysOut, days)
|
|
require.Equal(t, []timeRange{{from, to}}, repo.byDayCalls)
|
|
|
|
dailyOut, err := svc.UserDaily(context.Background(), uid, from, to)
|
|
require.NoError(t, err)
|
|
require.Equal(t, repo.daysOut, dailyOut)
|
|
require.Equal(t, []userTimeRange{{uid, from, to}}, repo.dailyCalls)
|
|
}
|