624 lines
21 KiB
Go
624 lines
21 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"strconv"
|
|
"strings"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/kirocooldown"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
type kiroUsageCooldownStore struct {
|
|
state *kirocooldown.State
|
|
err error
|
|
}
|
|
|
|
func (s *kiroUsageCooldownStore) ReserveRequest(context.Context, string) (time.Duration, error) {
|
|
return 0, nil
|
|
}
|
|
|
|
func (s *kiroUsageCooldownStore) MarkSuccess(context.Context, string) error {
|
|
return nil
|
|
}
|
|
|
|
func (s *kiroUsageCooldownStore) Mark429(context.Context, string) (time.Duration, error) {
|
|
return 0, nil
|
|
}
|
|
|
|
func (s *kiroUsageCooldownStore) MarkSuspended(context.Context, string) (time.Duration, error) {
|
|
return 0, nil
|
|
}
|
|
|
|
func (s *kiroUsageCooldownStore) GetState(context.Context, string) (*kirocooldown.State, error) {
|
|
return s.state, s.err
|
|
}
|
|
|
|
func (s *kiroUsageCooldownStore) ClearEarliestTransientCooldown(context.Context, []string) (bool, error) {
|
|
return false, nil
|
|
}
|
|
|
|
func kiroFloatPtr(v float64) *float64 {
|
|
return &v
|
|
}
|
|
|
|
func TestChannel_IsWebSearchEmulationEnabled_Kiro(t *testing.T) {
|
|
c := &Channel{
|
|
FeaturesConfig: map[string]any{
|
|
featureKeyWebSearchEmulation: map[string]any{"kiro": true},
|
|
},
|
|
}
|
|
|
|
require.True(t, c.IsWebSearchEmulationEnabled("kiro"))
|
|
}
|
|
|
|
func TestOpenAIGatewayServiceRecordUsage_NormalizesKiroBillingModel(t *testing.T) {
|
|
usageRepo := &openAIRecordUsageLogRepoStub{inserted: true}
|
|
userRepo := &openAIRecordUsageUserRepoStub{}
|
|
subRepo := &openAIRecordUsageSubRepoStub{}
|
|
svc := newOpenAIRecordUsageServiceForTest(usageRepo, userRepo, subRepo, nil)
|
|
svc.billingService = NewBillingService(svc.cfg, &PricingService{
|
|
pricingData: map[string]*LiteLLMModelPricing{
|
|
"claude-sonnet-4-6": {
|
|
InputCostPerToken: 2.5e-6,
|
|
OutputCostPerToken: 10e-6,
|
|
},
|
|
},
|
|
})
|
|
|
|
expectedCost, err := svc.billingService.CalculateCost("claude-sonnet-4-6", UsageTokens{
|
|
InputTokens: 20,
|
|
OutputTokens: 10,
|
|
}, 1.1)
|
|
require.NoError(t, err)
|
|
|
|
err = svc.RecordUsage(context.Background(), &OpenAIRecordUsageInput{
|
|
Result: &OpenAIForwardResult{
|
|
RequestID: "resp_kiro_billing_normalized",
|
|
Model: "claude-sonnet-4-6",
|
|
UpstreamModel: "claude-sonnet-4.6",
|
|
Usage: OpenAIUsage{
|
|
InputTokens: 20,
|
|
OutputTokens: 10,
|
|
},
|
|
Duration: time.Second,
|
|
},
|
|
APIKey: &APIKey{ID: 10},
|
|
User: &User{ID: 20},
|
|
Account: &Account{ID: 30, Platform: PlatformKiro},
|
|
})
|
|
|
|
require.NoError(t, err)
|
|
require.NotNil(t, usageRepo.lastLog)
|
|
require.Equal(t, "claude-sonnet-4-6", usageRepo.lastLog.Model)
|
|
require.Equal(t, "claude-sonnet-4-6", usageRepo.lastLog.RequestedModel)
|
|
require.NotNil(t, usageRepo.lastLog.UpstreamModel)
|
|
require.Equal(t, "claude-sonnet-4.6", *usageRepo.lastLog.UpstreamModel)
|
|
require.InDelta(t, expectedCost.ActualCost, usageRepo.lastLog.ActualCost, 1e-12)
|
|
require.InDelta(t, expectedCost.TotalCost, usageRepo.lastLog.TotalCost, 1e-12)
|
|
}
|
|
|
|
func TestAccountUsageService_GetUsage_KiroMapsCredits(t *testing.T) {
|
|
account := Account{
|
|
ID: 701,
|
|
Platform: PlatformKiro,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
"access_token": "kiro-access-token",
|
|
"provider": "Github",
|
|
"auth_method": "social",
|
|
"profile_arn": "arn:aws:codewhisperer:us-east-1:123456789012:profile/SOCIAL",
|
|
},
|
|
}
|
|
repo := &stubOpenAIAccountRepo{accounts: []Account{account}}
|
|
svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil)
|
|
|
|
resetAt := time.Now().Add(10 * 24 * time.Hour).Unix()
|
|
bonusExpiry := time.Now().Add(7 * 24 * time.Hour).Unix()
|
|
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
require.Equal(t, "/getUsageLimits", r.URL.Path)
|
|
require.Equal(t, "arn:aws:codewhisperer:us-east-1:123456789012:profile/SOCIAL", r.URL.Query().Get("profileArn"))
|
|
require.Equal(t, kiroUsageOrigin, r.URL.Query().Get("origin"))
|
|
require.Equal(t, kiroUsageResourceType, r.URL.Query().Get("resourceType"))
|
|
require.Equal(t, "Bearer kiro-access-token", r.Header.Get("Authorization"))
|
|
require.Equal(t, "*/*", r.Header.Get("Accept"))
|
|
require.True(t, strings.Contains(r.Header.Get("User-Agent"), "KiroIDE-"))
|
|
require.True(t, strings.Contains(r.Header.Get("X-Amz-User-Agent"), "KiroIDE-"))
|
|
require.Equal(t, "vibe", r.Header.Get("x-amzn-kiro-agent-mode"))
|
|
require.Equal(t, "true", r.Header.Get("x-amzn-codewhisperer-optout"))
|
|
require.NotEmpty(t, r.Header.Get("Amz-Sdk-Invocation-Id"))
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{
|
|
"nextDateReset": ` + strconv.FormatInt(resetAt, 10) + `,
|
|
"overageConfiguration": {"overageStatus":"ENABLED"},
|
|
"subscriptionInfo": {"subscriptionTitle":"KIRO PRO+","type":"Q_DEVELOPER_STANDALONE_PRO_PLUS"},
|
|
"usageBreakdownList": [{
|
|
"currency":"USD",
|
|
"currentOveragesWithPrecision":2,
|
|
"currentUsageWithPrecision":125,
|
|
"freeTrialInfo":{
|
|
"currentUsageWithPrecision":25,
|
|
"freeTrialExpiry":` + strconv.FormatInt(bonusExpiry, 10) + `,
|
|
"freeTrialStatus":"ACTIVE",
|
|
"usageLimitWithPrecision":500
|
|
},
|
|
"nextDateReset": ` + strconv.FormatInt(resetAt, 10) + `,
|
|
"overageCharges":0.08,
|
|
"resourceType":"CREDIT",
|
|
"usageLimitWithPrecision":2000
|
|
}]
|
|
}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
prevResolver := resolveKiroRuntimeEndpoint
|
|
resolveKiroRuntimeEndpoint = func(_ string) string { return server.URL }
|
|
defer func() { resolveKiroRuntimeEndpoint = prevResolver }()
|
|
|
|
usage, err := svc.GetUsage(context.Background(), account.ID)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, usage)
|
|
require.Equal(t, "active", usage.Source)
|
|
require.Equal(t, "KIRO PRO+", usage.KiroSubscriptionName)
|
|
require.Equal(t, "Q_DEVELOPER_STANDALONE_PRO_PLUS", usage.KiroSubscriptionType)
|
|
require.True(t, usage.KiroOveragesEnabled)
|
|
require.NotNil(t, usage.KiroCredit)
|
|
require.Equal(t, 125.0, usage.KiroCredit.CurrentUsage)
|
|
require.Equal(t, 2000.0, usage.KiroCredit.UsageLimit)
|
|
require.InDelta(t, 6.25, usage.KiroCredit.PercentageUsed, 0.001)
|
|
require.NotNil(t, usage.KiroBonus)
|
|
require.Equal(t, 25.0, usage.KiroBonus.CurrentUsage)
|
|
require.Equal(t, 500.0, usage.KiroBonus.UsageLimit)
|
|
require.NotNil(t, usage.KiroOverage)
|
|
require.Equal(t, "$", usage.KiroOverage.CurrencySymbol)
|
|
require.Equal(t, 2.0, usage.KiroOverage.CurrentOverages)
|
|
require.Equal(t, 0.08, usage.KiroOverage.OverageCharges)
|
|
require.NotNil(t, usage.KiroResetAt)
|
|
require.Equal(t, kiroQuotaStateOverageActive, usage.KiroQuotaState)
|
|
require.Equal(t, "overages_enabled", usage.KiroQuotaReason)
|
|
require.NotNil(t, usage.KiroQuotaResetAt)
|
|
}
|
|
|
|
func TestAccountUsageService_GetUsage_KiroActiveUsesCachedSnapshotWithinTTL(t *testing.T) {
|
|
account := Account{
|
|
ID: 702,
|
|
Platform: PlatformKiro,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
"access_token": "kiro-access-token",
|
|
"provider": "Github",
|
|
"auth_method": "social",
|
|
},
|
|
}
|
|
repo := &stubOpenAIAccountRepo{accounts: []Account{account}}
|
|
svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil)
|
|
|
|
successServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{
|
|
"subscriptionInfo": {"subscriptionTitle":"KIRO PRO+"},
|
|
"usageBreakdownList": [{
|
|
"currency":"USD",
|
|
"currentUsageWithPrecision":300,
|
|
"usageLimitWithPrecision":2000,
|
|
"resourceType":"CREDIT"
|
|
}]
|
|
}`))
|
|
}))
|
|
defer successServer.Close()
|
|
|
|
prevResolver := resolveKiroRuntimeEndpoint
|
|
resolveKiroRuntimeEndpoint = func(_ string) string { return successServer.URL }
|
|
defer func() { resolveKiroRuntimeEndpoint = prevResolver }()
|
|
|
|
firstUsage, err := svc.GetUsage(context.Background(), account.ID)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, firstUsage)
|
|
require.NotNil(t, firstUsage.KiroCredit)
|
|
require.Equal(t, 300.0, firstUsage.KiroCredit.CurrentUsage)
|
|
|
|
failingServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
http.Error(w, `{"message":"temporary failure"}`, http.StatusInternalServerError)
|
|
}))
|
|
defer failingServer.Close()
|
|
resolveKiroRuntimeEndpoint = func(_ string) string { return failingServer.URL }
|
|
|
|
usage, err := svc.GetUsage(context.Background(), account.ID)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, usage)
|
|
require.NotNil(t, usage.KiroCredit)
|
|
require.Equal(t, 300.0, usage.KiroCredit.CurrentUsage)
|
|
require.Empty(t, usage.Error)
|
|
require.Empty(t, usage.ErrorCode)
|
|
}
|
|
|
|
func TestAccountUsageService_GetUsage_KiroBuilderIDWithoutProfileArnOmitsProfileArn(t *testing.T) {
|
|
account := Account{
|
|
ID: 703,
|
|
Platform: PlatformKiro,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
"access_token": "kiro-access-token",
|
|
"provider": "BuilderId",
|
|
"auth_method": "idc",
|
|
"region": "us-east-1",
|
|
},
|
|
}
|
|
repo := &stubOpenAIAccountRepo{accounts: []Account{account}}
|
|
svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil)
|
|
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
require.Equal(t, "/getUsageLimits", r.URL.Path)
|
|
require.Empty(t, r.URL.Query().Get("profileArn"))
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{
|
|
"subscriptionInfo": {"subscriptionTitle":"KIRO PRO+"},
|
|
"usageBreakdownList": [{
|
|
"currency":"USD",
|
|
"currentUsageWithPrecision":42,
|
|
"usageLimitWithPrecision":2000,
|
|
"resourceType":"CREDIT"
|
|
}]
|
|
}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
prevResolver := resolveKiroRuntimeEndpoint
|
|
resolveKiroRuntimeEndpoint = func(_ string) string { return server.URL }
|
|
defer func() { resolveKiroRuntimeEndpoint = prevResolver }()
|
|
|
|
usage, err := svc.GetUsage(context.Background(), account.ID)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, usage)
|
|
require.NotNil(t, usage.KiroCredit)
|
|
require.Equal(t, 42.0, usage.KiroCredit.CurrentUsage)
|
|
}
|
|
|
|
func TestAccountUsageService_GetUsage_KiroEnterpriseUsesCredentialProfileArn(t *testing.T) {
|
|
account := Account{
|
|
ID: 707,
|
|
Platform: PlatformKiro,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
"access_token": "kiro-access-token",
|
|
"provider": "AWS",
|
|
"auth_method": "idc",
|
|
"region": "us-east-1",
|
|
"start_url": "https://d-example.awsapps.com/start",
|
|
"profile_arn": "arn:aws:codewhisperer:us-east-1:123456789012:profile/REALENTERPRISE",
|
|
},
|
|
}
|
|
repo := &stubOpenAIAccountRepo{accounts: []Account{account}}
|
|
svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil)
|
|
|
|
const resolvedProfileArn = "arn:aws:codewhisperer:us-east-1:123456789012:profile/REALENTERPRISE"
|
|
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
require.Equal(t, "/getUsageLimits", r.URL.Path)
|
|
require.Equal(t, resolvedProfileArn, r.URL.Query().Get("profileArn"))
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{
|
|
"subscriptionInfo": {"subscriptionTitle":"KIRO PRO+"},
|
|
"usageBreakdownList": [{
|
|
"currency":"USD",
|
|
"currentUsageWithPrecision":64,
|
|
"usageLimitWithPrecision":2000,
|
|
"resourceType":"CREDIT"
|
|
}]
|
|
}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
prevResolver := resolveKiroRuntimeEndpoint
|
|
resolveKiroRuntimeEndpoint = func(_ string) string { return server.URL }
|
|
defer func() { resolveKiroRuntimeEndpoint = prevResolver }()
|
|
|
|
usage, err := svc.GetUsage(context.Background(), account.ID)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, usage)
|
|
require.NotNil(t, usage.KiroCredit)
|
|
require.Equal(t, 64.0, usage.KiroCredit.CurrentUsage)
|
|
}
|
|
|
|
func TestAccountUsageService_GetUsage_KiroUsesAPIRegionForUsageRequest(t *testing.T) {
|
|
account := Account{
|
|
ID: 709,
|
|
Platform: PlatformKiro,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
"access_token": "kiro-access-token",
|
|
"provider": "AWS",
|
|
"auth_method": "idc",
|
|
"api_region": "eu-west-1",
|
|
"region": "ap-northeast-2",
|
|
"start_url": "https://d-example.awsapps.com/start",
|
|
"profile_arn": "arn:aws:codewhisperer:eu-west-1:123456789012:profile/REALAPIREGION",
|
|
},
|
|
}
|
|
repo := &stubOpenAIAccountRepo{accounts: []Account{account}}
|
|
svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil)
|
|
|
|
const resolvedProfileArn = "arn:aws:codewhisperer:eu-west-1:123456789012:profile/REALAPIREGION"
|
|
gotRegions := make([]string, 0, 2)
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
require.Equal(t, "/getUsageLimits", r.URL.Path)
|
|
require.Equal(t, resolvedProfileArn, r.URL.Query().Get("profileArn"))
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{
|
|
"subscriptionInfo": {"subscriptionTitle":"KIRO PRO+"},
|
|
"usageBreakdownList": [{
|
|
"currency":"USD",
|
|
"currentUsageWithPrecision":11,
|
|
"usageLimitWithPrecision":2000,
|
|
"resourceType":"CREDIT"
|
|
}]
|
|
}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
prevResolver := resolveKiroRuntimeEndpoint
|
|
resolveKiroRuntimeEndpoint = func(region string) string {
|
|
gotRegions = append(gotRegions, region)
|
|
return server.URL
|
|
}
|
|
defer func() { resolveKiroRuntimeEndpoint = prevResolver }()
|
|
|
|
usage, err := svc.GetUsage(context.Background(), account.ID)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, usage)
|
|
require.Equal(t, []string{"eu-west-1"}, gotRegions)
|
|
}
|
|
|
|
func TestAccountUsageService_GetUsage_KiroOmitsProfileArnAndUsesDefaultRegionWithoutAPIRegionOrProfileArn(t *testing.T) {
|
|
account := Account{
|
|
ID: 710,
|
|
Platform: PlatformKiro,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
"access_token": "kiro-access-token",
|
|
"provider": "AWS",
|
|
"auth_method": "idc",
|
|
"region": "ap-northeast-2",
|
|
"start_url": "https://d-example.awsapps.com/start",
|
|
},
|
|
}
|
|
repo := &stubOpenAIAccountRepo{accounts: []Account{account}}
|
|
svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil)
|
|
|
|
gotRegions := make([]string, 0, 2)
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
require.Equal(t, "/getUsageLimits", r.URL.Path)
|
|
require.Empty(t, r.URL.Query().Get("profileArn"))
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{
|
|
"subscriptionInfo": {"subscriptionTitle":"KIRO PRO+"},
|
|
"usageBreakdownList": [{
|
|
"currency":"USD",
|
|
"currentUsageWithPrecision":7,
|
|
"usageLimitWithPrecision":2000,
|
|
"resourceType":"CREDIT"
|
|
}]
|
|
}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
prevResolver := resolveKiroRuntimeEndpoint
|
|
resolveKiroRuntimeEndpoint = func(region string) string {
|
|
gotRegions = append(gotRegions, region)
|
|
return server.URL
|
|
}
|
|
defer func() { resolveKiroRuntimeEndpoint = prevResolver }()
|
|
|
|
usage, err := svc.GetUsage(context.Background(), account.ID)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, usage)
|
|
require.Equal(t, []string{kiroDefaultRegion}, gotRegions)
|
|
}
|
|
|
|
func TestAccountUsageService_GetUsage_KiroIncludesRuntimeCooldownState(t *testing.T) {
|
|
account := Account{
|
|
ID: 704,
|
|
Platform: PlatformKiro,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
"access_token": "kiro-access-token",
|
|
"provider": "Github",
|
|
"auth_method": "social",
|
|
},
|
|
}
|
|
repo := &stubOpenAIAccountRepo{accounts: []Account{account}}
|
|
svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil).
|
|
SetKiroCooldownStore(&kiroUsageCooldownStore{
|
|
state: &kirocooldown.State{
|
|
Active: true,
|
|
Reason: kirocooldown.CooldownReason429,
|
|
CooldownUntil: time.Now().Add(90 * time.Second),
|
|
Remaining: 90 * time.Second,
|
|
},
|
|
})
|
|
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{
|
|
"subscriptionInfo": {"subscriptionTitle":"KIRO PRO+"},
|
|
"usageBreakdownList": [{
|
|
"currency":"USD",
|
|
"currentUsageWithPrecision":42,
|
|
"usageLimitWithPrecision":2000,
|
|
"resourceType":"CREDIT"
|
|
}]
|
|
}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
prevResolver := resolveKiroRuntimeEndpoint
|
|
resolveKiroRuntimeEndpoint = func(_ string) string { return server.URL }
|
|
defer func() { resolveKiroRuntimeEndpoint = prevResolver }()
|
|
|
|
usage, err := svc.GetUsage(context.Background(), account.ID)
|
|
require.NoError(t, err)
|
|
require.Equal(t, "cooldown", usage.KiroRuntimeState)
|
|
require.Equal(t, kirocooldown.CooldownReason429, usage.KiroRuntimeReason)
|
|
require.NotNil(t, usage.KiroRuntimeResetAt)
|
|
}
|
|
|
|
func TestBuildKiroDegradedUsage_ClassifiesProfileError(t *testing.T) {
|
|
info := buildKiroDegradedUsage(&kiroUsageHTTPError{
|
|
StatusCode: http.StatusBadRequest,
|
|
Body: `{"message":"profileArn is required for this request."}`,
|
|
})
|
|
|
|
require.Equal(t, errorCodeForbidden, info.ErrorCode)
|
|
require.False(t, info.NeedsReauth)
|
|
}
|
|
|
|
func TestBuildKiroDegradedUsage_ClassifiesOverageExhausted(t *testing.T) {
|
|
info := buildKiroDegradedUsage(&kiroUsageHTTPError{
|
|
StatusCode: http.StatusTooManyRequests,
|
|
Body: `{"message":"overage exhausted for this billing window"}`,
|
|
})
|
|
|
|
require.Equal(t, errorCodeNetworkError, info.ErrorCode)
|
|
require.Equal(t, kiroQuotaStateOverageExhausted, info.KiroQuotaState)
|
|
require.Contains(t, info.KiroQuotaReason, "overage exhausted")
|
|
}
|
|
|
|
func TestAccountUsageService_GetUsage_KiroCachesErrorSnapshotWhenRefreshFailsWithoutPriorSuccess(t *testing.T) {
|
|
account := Account{
|
|
ID: 708,
|
|
Platform: PlatformKiro,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
"access_token": "kiro-access-token",
|
|
"provider": "Github",
|
|
"auth_method": "social",
|
|
},
|
|
}
|
|
repo := &stubOpenAIAccountRepo{accounts: []Account{account}}
|
|
svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil)
|
|
|
|
requestCount := 0
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
requestCount++
|
|
http.Error(w, `{"message":"FEATURE_NOT_SUPPORTED","reason":"FEATURE_NOT_SUPPORTED"}`, http.StatusForbidden)
|
|
}))
|
|
defer server.Close()
|
|
|
|
prevResolver := resolveKiroRuntimeEndpoint
|
|
resolveKiroRuntimeEndpoint = func(_ string) string { return server.URL }
|
|
defer func() { resolveKiroRuntimeEndpoint = prevResolver }()
|
|
|
|
firstUsage, err := svc.GetUsage(context.Background(), account.ID)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, firstUsage)
|
|
require.Equal(t, errorCodeForbidden, firstUsage.ErrorCode)
|
|
|
|
secondUsage, err := svc.GetUsage(context.Background(), account.ID)
|
|
require.NoError(t, err)
|
|
require.NotNil(t, secondUsage)
|
|
require.Equal(t, errorCodeForbidden, secondUsage.ErrorCode)
|
|
require.Equal(t, 1, requestCount)
|
|
}
|
|
|
|
func TestMapKiroUsageToInfo_CreditsExhaustedWithoutOverages(t *testing.T) {
|
|
info := mapKiroUsageToInfo(&kiroUsageLimitsResponse{
|
|
NextDateReset: "2099-03-13T12:00:00Z",
|
|
OverageConfiguration: kiroOverageConfiguration{
|
|
OverageStatus: "DISABLED",
|
|
},
|
|
UsageBreakdownList: []kiroUsageBreakdown{
|
|
{
|
|
ResourceType: "CREDIT",
|
|
CurrentUsageWithPrecision: kiroFloatPtr(2000),
|
|
UsageLimitWithPrecision: kiroFloatPtr(2000),
|
|
CurrentOveragesWithPrecision: kiroFloatPtr(0),
|
|
},
|
|
},
|
|
})
|
|
|
|
require.Equal(t, kiroQuotaStateCreditsExhausted, info.KiroQuotaState)
|
|
require.Equal(t, "credits_exhausted", info.KiroQuotaReason)
|
|
require.NotNil(t, info.KiroQuotaResetAt)
|
|
}
|
|
|
|
func TestAccountUsageService_EnrichAccountWithKiroRuntimeState(t *testing.T) {
|
|
svc := NewAccountUsageService(nil, nil, nil, nil, nil, NewUsageCache(), nil, nil).
|
|
SetKiroCooldownStore(&kiroUsageCooldownStore{
|
|
state: &kirocooldown.State{
|
|
Active: true,
|
|
Reason: kirocooldown.CooldownReason429,
|
|
CooldownUntil: time.Now().Add(2 * time.Minute),
|
|
Remaining: 2 * time.Minute,
|
|
},
|
|
})
|
|
|
|
account := &Account{
|
|
ID: 705,
|
|
Platform: PlatformKiro,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{"access_token": "kiro-access-token"},
|
|
}
|
|
|
|
svc.EnrichAccountWithKiroRuntimeState(context.Background(), account)
|
|
require.Equal(t, "cooldown", account.KiroRuntimeState)
|
|
require.Equal(t, kirocooldown.CooldownReason429, account.KiroRuntimeReason)
|
|
require.NotNil(t, account.KiroRuntimeResetAt)
|
|
}
|
|
|
|
func TestAccountUsageService_EnrichAccountWithKiroRuntimeStateIncludesCachedQuotaState(t *testing.T) {
|
|
account := Account{
|
|
ID: 706,
|
|
Platform: PlatformKiro,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{
|
|
"access_token": "kiro-access-token",
|
|
"provider": "Github",
|
|
"auth_method": "social",
|
|
},
|
|
}
|
|
repo := &stubOpenAIAccountRepo{accounts: []Account{account}}
|
|
svc := NewAccountUsageService(repo, nil, nil, nil, nil, NewUsageCache(), nil, nil)
|
|
|
|
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
w.Header().Set("Content-Type", "application/json")
|
|
_, _ = w.Write([]byte(`{
|
|
"nextDateReset":"2099-03-13T12:00:00Z",
|
|
"overageConfiguration":{"overageStatus":"ENABLED"},
|
|
"subscriptionInfo": {"subscriptionTitle":"KIRO PRO+"},
|
|
"usageBreakdownList": [{
|
|
"currency":"USD",
|
|
"currentUsageWithPrecision":2000,
|
|
"currentOveragesWithPrecision":4,
|
|
"overageCharges":0.2,
|
|
"usageLimitWithPrecision":2000,
|
|
"resourceType":"CREDIT"
|
|
}]
|
|
}`))
|
|
}))
|
|
defer server.Close()
|
|
|
|
prevResolver := resolveKiroRuntimeEndpoint
|
|
resolveKiroRuntimeEndpoint = func(_ string) string { return server.URL }
|
|
defer func() { resolveKiroRuntimeEndpoint = prevResolver }()
|
|
|
|
_, err := svc.GetUsage(context.Background(), account.ID)
|
|
require.NoError(t, err)
|
|
|
|
target := &Account{
|
|
ID: account.ID,
|
|
Platform: PlatformKiro,
|
|
Type: AccountTypeOAuth,
|
|
Credentials: map[string]any{"access_token": "kiro-access-token"},
|
|
}
|
|
svc.EnrichAccountWithKiroRuntimeState(context.Background(), target)
|
|
|
|
require.Equal(t, kiroQuotaStateOverageActive, target.KiroQuotaState)
|
|
require.Equal(t, "overages_enabled", target.KiroQuotaReason)
|
|
require.NotNil(t, target.KiroQuotaResetAt)
|
|
}
|