312 lines
10 KiB
Go
312 lines
10 KiB
Go
|
package handler
|
||
|
|
|
||
|
|
import (
|
||
|
|
"context"
|
||
|
|
"encoding/json"
|
||
|
|
"errors"
|
||
|
|
"net/http"
|
||
|
|
"net/http/httptest"
|
||
|
|
"testing"
|
||
|
|
"time"
|
||
|
|
|
||
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/timezone"
|
||
|
|
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||
|
|
"github.com/Wei-Shaw/sub2api/internal/service"
|
||
|
|
"github.com/gin-gonic/gin"
|
||
|
|
"github.com/stretchr/testify/require"
|
||
|
|
)
|
||
|
|
|
||
|
|
type keyBillingUserGroupRateRepo struct {
|
||
|
|
service.UserGroupRateRepository
|
||
|
|
rate *float64
|
||
|
|
err error
|
||
|
|
gotUserID int64
|
||
|
|
gotGroupID int64
|
||
|
|
lookupCalls int
|
||
|
|
}
|
||
|
|
|
||
|
|
func (r *keyBillingUserGroupRateRepo) GetByUserAndGroup(_ context.Context, userID, groupID int64) (*float64, error) {
|
||
|
|
r.gotUserID = userID
|
||
|
|
r.gotGroupID = groupID
|
||
|
|
r.lookupCalls++
|
||
|
|
return r.rate, r.err
|
||
|
|
}
|
||
|
|
|
||
|
|
func newKeyBillingHandler(repo service.UserGroupRateRepository) *GatewayHandler {
|
||
|
|
return &GatewayHandler{
|
||
|
|
gatewayService: newKeyBillingGatewayService(repo),
|
||
|
|
openAIGatewayService: newKeyBillingOpenAIGatewayService(repo),
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func newKeyBillingGatewayService(repo service.UserGroupRateRepository) *service.GatewayService {
|
||
|
|
return service.NewGatewayService(
|
||
|
|
nil, nil, nil, nil, nil, nil, repo, nil, nil, nil, nil, nil, nil, nil,
|
||
|
|
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||
|
|
)
|
||
|
|
}
|
||
|
|
|
||
|
|
func newKeyBillingOpenAIGatewayService(repo service.UserGroupRateRepository) *service.OpenAIGatewayService {
|
||
|
|
return service.NewOpenAIGatewayService(
|
||
|
|
nil, nil, nil, nil, nil, repo, nil, nil, nil, nil, nil,
|
||
|
|
nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil,
|
||
|
|
)
|
||
|
|
}
|
||
|
|
|
||
|
|
func newKeyBillingContext(apiKey *service.APIKey) (*gin.Context, *httptest.ResponseRecorder) {
|
||
|
|
gin.SetMode(gin.TestMode)
|
||
|
|
w := httptest.NewRecorder()
|
||
|
|
c, _ := gin.CreateTestContext(w)
|
||
|
|
c.Request = httptest.NewRequest(http.MethodGet, "/v1/sub2api/billing", nil)
|
||
|
|
if apiKey != nil {
|
||
|
|
c.Set(string(middleware2.ContextKeyAPIKey), apiKey)
|
||
|
|
}
|
||
|
|
return c, w
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestGatewayHandlerKeyBillingInfoUsesGroupRate(t *testing.T) {
|
||
|
|
groupID := int64(7)
|
||
|
|
apiKey := &service.APIKey{
|
||
|
|
UserID: 11,
|
||
|
|
GroupID: &groupID,
|
||
|
|
Key: "sk-sensitive-value",
|
||
|
|
Group: &service.Group{
|
||
|
|
ID: groupID,
|
||
|
|
Name: "private-group-name",
|
||
|
|
RateMultiplier: 0.75,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
c, w := newKeyBillingContext(apiKey)
|
||
|
|
|
||
|
|
newKeyBillingHandler(nil).KeyBillingInfo(c)
|
||
|
|
|
||
|
|
require.Equal(t, http.StatusOK, w.Code)
|
||
|
|
require.Equal(t, "no-store", w.Header().Get("Cache-Control"))
|
||
|
|
var got keyBillingInfoResponse
|
||
|
|
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got))
|
||
|
|
require.Equal(t, "sub2api.key_billing", got.Object)
|
||
|
|
require.Equal(t, 1, got.SchemaVersion)
|
||
|
|
require.Equal(t, "token", got.BillingScope)
|
||
|
|
require.Equal(t, 0.75, got.GroupRateMultiplier)
|
||
|
|
require.Nil(t, got.UserRateMultiplier)
|
||
|
|
require.Equal(t, 0.75, got.ResolvedRateMultiplier)
|
||
|
|
require.False(t, got.PeakRateEnabled)
|
||
|
|
require.Nil(t, got.PeakStart)
|
||
|
|
require.Nil(t, got.PeakEnd)
|
||
|
|
require.Nil(t, got.PeakRateMultiplier)
|
||
|
|
require.Nil(t, got.AppliedPeakMultiplier)
|
||
|
|
require.Equal(t, 0.75, got.EffectiveRateMultiplier)
|
||
|
|
require.Nil(t, got.Timezone)
|
||
|
|
require.False(t, got.ObservedAt.IsZero())
|
||
|
|
var fields map[string]json.RawMessage
|
||
|
|
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &fields))
|
||
|
|
require.NotContains(t, fields, "user_rate_multiplier")
|
||
|
|
require.NotContains(t, fields, "peak_start")
|
||
|
|
require.NotContains(t, fields, "peak_end")
|
||
|
|
require.NotContains(t, fields, "peak_rate_multiplier")
|
||
|
|
require.NotContains(t, fields, "applied_peak_multiplier")
|
||
|
|
require.NotContains(t, fields, "timezone")
|
||
|
|
require.NotContains(t, w.Body.String(), apiKey.Key)
|
||
|
|
require.NotContains(t, w.Body.String(), apiKey.Group.Name)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestGatewayHandlerKeyBillingInfoUsesUserOverride(t *testing.T) {
|
||
|
|
groupID := int64(7)
|
||
|
|
userRate := 0.5
|
||
|
|
apiKey := &service.APIKey{
|
||
|
|
UserID: 11,
|
||
|
|
GroupID: &groupID,
|
||
|
|
Group: &service.Group{ID: groupID, RateMultiplier: 0.75},
|
||
|
|
}
|
||
|
|
c, w := newKeyBillingContext(apiKey)
|
||
|
|
repo := &keyBillingUserGroupRateRepo{rate: &userRate}
|
||
|
|
|
||
|
|
newKeyBillingHandler(repo).KeyBillingInfo(c)
|
||
|
|
|
||
|
|
require.Equal(t, http.StatusOK, w.Code)
|
||
|
|
require.Equal(t, 1, repo.lookupCalls)
|
||
|
|
require.Equal(t, apiKey.UserID, repo.gotUserID)
|
||
|
|
require.Equal(t, groupID, repo.gotGroupID)
|
||
|
|
var got keyBillingInfoResponse
|
||
|
|
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got))
|
||
|
|
require.NotNil(t, got.UserRateMultiplier)
|
||
|
|
require.Equal(t, 0.5, *got.UserRateMultiplier)
|
||
|
|
require.Equal(t, 0.5, got.ResolvedRateMultiplier)
|
||
|
|
require.Equal(t, 0.5, got.EffectiveRateMultiplier)
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestBuildKeyBillingInfoAppliesPeakMultiplier(t *testing.T) {
|
||
|
|
groupID := int64(7)
|
||
|
|
apiKey := &service.APIKey{
|
||
|
|
GroupID: &groupID,
|
||
|
|
Group: &service.Group{
|
||
|
|
ID: groupID,
|
||
|
|
RateMultiplier: 1.2,
|
||
|
|
SubscriptionType: service.SubscriptionTypeSubscription,
|
||
|
|
PeakRateEnabled: true,
|
||
|
|
PeakStart: "09:00",
|
||
|
|
PeakEnd: "18:00",
|
||
|
|
PeakRateMultiplier: 1.5,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
now := time.Date(2026, time.July, 12, 10, 0, 0, 0, timezone.Location())
|
||
|
|
userRate := 0.8
|
||
|
|
|
||
|
|
got := buildKeyBillingInfo(apiKey, userRate, now)
|
||
|
|
|
||
|
|
require.Equal(t, 1.2, got.GroupRateMultiplier)
|
||
|
|
require.NotNil(t, got.UserRateMultiplier)
|
||
|
|
require.Equal(t, 0.8, *got.UserRateMultiplier)
|
||
|
|
require.Equal(t, 0.8, got.ResolvedRateMultiplier)
|
||
|
|
require.True(t, got.PeakRateEnabled)
|
||
|
|
require.NotNil(t, got.PeakStart)
|
||
|
|
require.Equal(t, "09:00", *got.PeakStart)
|
||
|
|
require.NotNil(t, got.PeakEnd)
|
||
|
|
require.Equal(t, "18:00", *got.PeakEnd)
|
||
|
|
require.NotNil(t, got.PeakRateMultiplier)
|
||
|
|
require.Equal(t, 1.5, *got.PeakRateMultiplier)
|
||
|
|
require.NotNil(t, got.AppliedPeakMultiplier)
|
||
|
|
require.Equal(t, 1.5, *got.AppliedPeakMultiplier)
|
||
|
|
require.InDelta(t, 1.2, got.EffectiveRateMultiplier, 1e-12)
|
||
|
|
require.NotNil(t, got.Timezone)
|
||
|
|
require.Equal(t, timezone.Location().String(), *got.Timezone)
|
||
|
|
require.Equal(t, now.UTC(), got.ObservedAt)
|
||
|
|
|
||
|
|
encoded, err := json.Marshal(got)
|
||
|
|
require.NoError(t, err)
|
||
|
|
var fields map[string]json.RawMessage
|
||
|
|
require.NoError(t, json.Unmarshal(encoded, &fields))
|
||
|
|
for _, field := range []string{
|
||
|
|
"user_rate_multiplier",
|
||
|
|
"peak_start",
|
||
|
|
"peak_end",
|
||
|
|
"peak_rate_multiplier",
|
||
|
|
"applied_peak_multiplier",
|
||
|
|
"timezone",
|
||
|
|
} {
|
||
|
|
require.Contains(t, fields, field)
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestKeyBillingInfoJSONKeepsZeroPeakMultiplierWhenEnabled(t *testing.T) {
|
||
|
|
groupID := int64(7)
|
||
|
|
apiKey := &service.APIKey{
|
||
|
|
GroupID: &groupID,
|
||
|
|
Group: &service.Group{
|
||
|
|
ID: groupID,
|
||
|
|
SubscriptionType: service.SubscriptionTypeSubscription,
|
||
|
|
PeakRateEnabled: true,
|
||
|
|
PeakStart: "00:00",
|
||
|
|
PeakEnd: "23:59",
|
||
|
|
PeakRateMultiplier: 0,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
now := time.Date(2026, time.July, 12, 12, 0, 0, 0, timezone.Location())
|
||
|
|
encoded, err := json.Marshal(buildKeyBillingInfo(apiKey, apiKey.Group.RateMultiplier, now))
|
||
|
|
require.NoError(t, err)
|
||
|
|
|
||
|
|
var fields map[string]json.RawMessage
|
||
|
|
require.NoError(t, json.Unmarshal(encoded, &fields))
|
||
|
|
require.JSONEq(t, "0", string(fields["peak_rate_multiplier"]))
|
||
|
|
require.JSONEq(t, "0", string(fields["applied_peak_multiplier"]))
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestGatewayHandlerKeyBillingInfoErrorsAreSafe(t *testing.T) {
|
||
|
|
t.Run("missing API key", func(t *testing.T) {
|
||
|
|
c, w := newKeyBillingContext(nil)
|
||
|
|
newKeyBillingHandler(nil).KeyBillingInfo(c)
|
||
|
|
require.Equal(t, http.StatusUnauthorized, w.Code)
|
||
|
|
})
|
||
|
|
|
||
|
|
t.Run("ungrouped API key", func(t *testing.T) {
|
||
|
|
c, w := newKeyBillingContext(&service.APIKey{})
|
||
|
|
newKeyBillingHandler(nil).KeyBillingInfo(c)
|
||
|
|
require.Equal(t, http.StatusForbidden, w.Code)
|
||
|
|
})
|
||
|
|
|
||
|
|
t.Run("missing billing service", func(t *testing.T) {
|
||
|
|
groupID := int64(7)
|
||
|
|
c, w := newKeyBillingContext(&service.APIKey{
|
||
|
|
UserID: 11,
|
||
|
|
GroupID: &groupID,
|
||
|
|
Group: &service.Group{ID: groupID, RateMultiplier: 1},
|
||
|
|
})
|
||
|
|
(&GatewayHandler{}).KeyBillingInfo(c)
|
||
|
|
require.Equal(t, http.StatusInternalServerError, w.Code)
|
||
|
|
})
|
||
|
|
|
||
|
|
t.Run("rate lookup failure matches billing fallback", func(t *testing.T) {
|
||
|
|
groupID := int64(7)
|
||
|
|
c, w := newKeyBillingContext(&service.APIKey{
|
||
|
|
UserID: 11,
|
||
|
|
GroupID: &groupID,
|
||
|
|
Group: &service.Group{ID: groupID, RateMultiplier: 1},
|
||
|
|
})
|
||
|
|
newKeyBillingHandler(&keyBillingUserGroupRateRepo{err: errors.New("database password leaked")}).KeyBillingInfo(c)
|
||
|
|
require.Equal(t, http.StatusOK, w.Code)
|
||
|
|
var got keyBillingInfoResponse
|
||
|
|
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got))
|
||
|
|
require.Equal(t, 1.0, got.ResolvedRateMultiplier)
|
||
|
|
require.NotContains(t, w.Body.String(), "database password leaked")
|
||
|
|
})
|
||
|
|
}
|
||
|
|
|
||
|
|
func TestGatewayHandlerKeyBillingInfoSharesBillingResolverCacheByPlatform(t *testing.T) {
|
||
|
|
for _, tc := range []struct {
|
||
|
|
name string
|
||
|
|
platform string
|
||
|
|
openAI bool
|
||
|
|
}{
|
||
|
|
{name: "anthropic", platform: service.PlatformAnthropic},
|
||
|
|
{name: "openai", platform: service.PlatformOpenAI, openAI: true},
|
||
|
|
{name: "grok", platform: service.PlatformGrok, openAI: true},
|
||
|
|
} {
|
||
|
|
t.Run(tc.name, func(t *testing.T) {
|
||
|
|
groupID := int64(7)
|
||
|
|
oldRate, newRate := 0.5, 1.8
|
||
|
|
repo := &keyBillingUserGroupRateRepo{rate: &oldRate}
|
||
|
|
gatewayService := newKeyBillingGatewayService(repo)
|
||
|
|
openAIGatewayService := newKeyBillingOpenAIGatewayService(repo)
|
||
|
|
h := &GatewayHandler{
|
||
|
|
gatewayService: gatewayService,
|
||
|
|
openAIGatewayService: openAIGatewayService,
|
||
|
|
}
|
||
|
|
apiKey := &service.APIKey{
|
||
|
|
UserID: 11,
|
||
|
|
GroupID: &groupID,
|
||
|
|
Group: &service.Group{
|
||
|
|
ID: groupID,
|
||
|
|
Platform: tc.platform,
|
||
|
|
RateMultiplier: 0.75,
|
||
|
|
},
|
||
|
|
}
|
||
|
|
|
||
|
|
if tc.openAI {
|
||
|
|
require.Equal(t, oldRate, openAIGatewayService.ResolveUserGroupRateMultiplier(context.Background(), apiKey.UserID, groupID, apiKey.Group.RateMultiplier))
|
||
|
|
} else {
|
||
|
|
require.Equal(t, oldRate, gatewayService.ResolveUserGroupRateMultiplier(context.Background(), apiKey.UserID, groupID, apiKey.Group.RateMultiplier))
|
||
|
|
}
|
||
|
|
repo.rate = &newRate
|
||
|
|
|
||
|
|
for range 2 {
|
||
|
|
c, w := newKeyBillingContext(apiKey)
|
||
|
|
h.KeyBillingInfo(c)
|
||
|
|
require.Equal(t, http.StatusOK, w.Code)
|
||
|
|
var got keyBillingInfoResponse
|
||
|
|
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &got))
|
||
|
|
require.Equal(t, oldRate, got.ResolvedRateMultiplier)
|
||
|
|
require.Equal(t, oldRate, got.EffectiveRateMultiplier)
|
||
|
|
}
|
||
|
|
|
||
|
|
var billedRate float64
|
||
|
|
if tc.openAI {
|
||
|
|
billedRate = openAIGatewayService.ResolveUserGroupRateMultiplier(context.Background(), apiKey.UserID, groupID, apiKey.Group.RateMultiplier)
|
||
|
|
} else {
|
||
|
|
billedRate = gatewayService.ResolveUserGroupRateMultiplier(context.Background(), apiKey.UserID, groupID, apiKey.Group.RateMultiplier)
|
||
|
|
}
|
||
|
|
require.Equal(t, oldRate, billedRate)
|
||
|
|
require.Equal(t, 1, repo.lookupCalls)
|
||
|
|
})
|
||
|
|
}
|
||
|
|
}
|