| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176 |
- package xai
- import (
- "encoding/json"
- "net/http"
- "testing"
- "github.com/stretchr/testify/require"
- )
- func TestBuildBillingURL(t *testing.T) {
- t.Parallel()
- require.Equal(t, "https://cli-chat-proxy.grok.com/v1/billing?format=credits", BuildBillingURL(true))
- require.Equal(t, "https://cli-chat-proxy.grok.com/v1/billing", BuildBillingURL(false))
- }
- func TestBuildBillingURLWithValidator(t *testing.T) {
- t.Parallel()
- weeklyURL, err := BuildBillingURLWithValidator(DefaultCLIBaseURL, true, ValidateTrustedBaseURL)
- require.NoError(t, err)
- require.Equal(t, "https://cli-chat-proxy.grok.com/v1/billing?format=credits", weeklyURL)
- monthlyURL, err := BuildBillingURLWithValidator("https://relay.example.test/v1", false, ValidateBaseURL)
- require.NoError(t, err)
- require.Equal(t, "https://relay.example.test/v1/billing", monthlyURL)
- _, err = BuildBillingURLWithValidator("https://relay.example.test/v1", true, ValidateTrustedBaseURL)
- require.Error(t, err)
- }
- func TestApplyCLIBillingHeaders(t *testing.T) {
- t.Parallel()
- req, err := http.NewRequest(http.MethodGet, BuildBillingURL(true), nil)
- require.NoError(t, err)
- ApplyCLIBillingHeaders(req, " token ")
- require.Equal(t, "Bearer token", req.Header.Get("Authorization"))
- require.Equal(t, CLITokenAuthValue, req.Header.Get(CLITokenAuthHeader))
- require.Equal(t, CLIClientVersion, req.Header.Get(CLIClientVersionHeader))
- require.Equal(t, "grok-pager/"+CLIClientVersion+" grok-shell/"+CLIClientVersion+" (macos; aarch64)", req.UserAgent())
- }
- func TestBuildBillingSummaryWeeklyAndMonthly(t *testing.T) {
- t.Parallel()
- weeklyBody := []byte(`{
- "config": {
- "currentPeriod": {"type":"WEEKLY","start":"2026-07-09T03:25:00Z","end":"2026-07-16T03:25:00Z"},
- "creditUsagePercent": 2.0,
- "productUsage": [{"product":"Api","usagePercent":2.0}],
- "prepaidBalance": {"val": 12},
- "onDemandCap": {"val": 100},
- "onDemandUsed": {"val": 5},
- "isUnifiedBillingUser": true
- }
- }`)
- monthlyBody := []byte(`{
- "config": {
- "monthlyLimit": {"val": 15000},
- "used": {"val": 78},
- "billingPeriodStart": "2026-07-01T00:00:00Z",
- "billingPeriodEnd": "2026-08-01T00:00:00Z"
- }
- }`)
- weeklyPayload, err := ParseBillingPayload(weeklyBody)
- require.NoError(t, err)
- monthlyPayload, err := ParseBillingPayload(monthlyBody)
- require.NoError(t, err)
- weekly := BuildBillingSummary(weeklyPayload.Config)
- monthly := BuildBillingSummary(monthlyPayload.Config)
- require.NotNil(t, weekly)
- require.NotNil(t, monthly)
- require.Equal(t, "weekly", weekly.PeriodType)
- require.InDelta(t, 2.0, *weekly.UsagePercent, 1e-9)
- require.Equal(t, "Api", weekly.ProductUsage[0].Product)
- require.InDelta(t, 12, *weekly.PrepaidBalance, 1e-9)
- require.InDelta(t, 100, *weekly.OnDemandCap, 1e-9)
- require.InDelta(t, 5, *weekly.OnDemandUsed, 1e-9)
- require.True(t, weekly.IsUnifiedBillingUser)
- require.Equal(t, "SuperGrok", monthly.Plan)
- require.InDelta(t, 15000, *monthly.MonthlyLimitCents, 1e-9)
- require.InDelta(t, 78, *monthly.UsedCents, 1e-9)
- require.InDelta(t, 0.52, *monthly.UsedPercent, 1e-2)
- require.InDelta(t, 150, *monthly.MonthlyLimit, 1e-9)
- require.InDelta(t, 0.78, *monthly.MonthlyUsed, 1e-9)
- merged := MergeBillingProbeResult(nil, weekly, monthly, true, true)
- require.Equal(t, "weekly", merged.PeriodType)
- require.InDelta(t, 2.0, *merged.UsagePercent, 1e-9)
- require.Equal(t, "SuperGrok", merged.Plan)
- require.InDelta(t, 15000, *merged.MonthlyLimitCents, 1e-9)
- require.Equal(t, "2026-08-01T00:00:00Z", merged.BillingPeriodEnd)
- require.InDelta(t, 12, *merged.PrepaidBalance, 1e-9)
- require.InDelta(t, 100, *merged.OnDemandCap, 1e-9)
- require.InDelta(t, 150, *merged.MonthlyLimit, 1e-9)
- require.InDelta(t, 0.78, *merged.MonthlyUsed, 1e-9)
- }
- func TestParseCentValueBareNumber(t *testing.T) {
- t.Parallel()
- raw, _ := json.Marshal(15000)
- v := parseCentValue(raw)
- require.NotNil(t, v)
- require.InDelta(t, 15000, *v, 1e-9)
- }
- func TestBuildBillingSummaryMonthlyOnlyKeepsWeeklyUsageEmpty(t *testing.T) {
- t.Parallel()
- payload, err := ParseBillingPayload([]byte(`{"config":{"monthlyLimit":{"val":15000},"used":{"val":7500},"billingPeriodStart":"2026-07-01T00:00:00Z","billingPeriodEnd":"2026-08-01T00:00:00Z"}}`))
- require.NoError(t, err)
- summary := BuildBillingSummary(payload.Config)
- require.NotNil(t, summary)
- require.Equal(t, "monthly", summary.PeriodType)
- require.Nil(t, summary.UsagePercent)
- require.InDelta(t, 50, *summary.UsedPercent, 1e-9)
- }
- func TestBuildBillingSummaryWeeklyDoesNotInheritMonthlyPeriodEnd(t *testing.T) {
- t.Parallel()
- // Weekly usage without currentPeriod.end must not copy billingPeriodEnd (monthly).
- payload, err := ParseBillingPayload([]byte(`{
- "config": {
- "creditUsagePercent": 95.0,
- "productUsage": [{"product":"Api","usagePercent":95.0}],
- "billingPeriodStart": "2026-07-01T00:00:00Z",
- "billingPeriodEnd": "2026-08-01T00:00:00Z",
- "monthlyLimit": {"val": 15000},
- "used": {"val": 1000}
- }
- }`))
- require.NoError(t, err)
- summary := BuildBillingSummary(payload.Config)
- require.NotNil(t, summary)
- require.Equal(t, "weekly", summary.PeriodType)
- require.Equal(t, "", summary.PeriodEnd)
- require.Equal(t, "2026-08-01T00:00:00Z", summary.BillingPeriodEnd)
- }
- func TestMergeBillingProbeResultRetainsFailedWindow(t *testing.T) {
- t.Parallel()
- previous := &BillingSummary{
- PeriodType: "weekly",
- UsagePercent: floatPointer(100),
- PeriodEnd: "2026-07-16T00:00:00Z",
- MonthlyLimitCents: floatPointer(15000),
- UsedPercent: floatPointer(20),
- BillingPeriodEnd: "2026-08-01T00:00:00Z",
- WeeklyUpdatedAt: "2026-07-10T00:00:00Z",
- MonthlyUpdatedAt: "2026-07-10T00:00:00Z",
- FailedWindows: []string{"monthly"},
- }
- monthly := &BillingSummary{
- PeriodType: "monthly",
- MonthlyLimitCents: floatPointer(15000),
- UsedPercent: floatPointer(30),
- BillingPeriodEnd: "2026-08-01T00:00:00Z",
- }
- merged := MergeBillingProbeResult(previous, nil, monthly, false, true)
- require.Equal(t, "weekly", merged.PeriodType)
- require.InDelta(t, 100, *merged.UsagePercent, 1e-9)
- require.Equal(t, previous.WeeklyUpdatedAt, merged.WeeklyUpdatedAt)
- require.InDelta(t, 30, *merged.UsedPercent, 1e-9)
- require.NotEqual(t, previous.MonthlyUpdatedAt, merged.MonthlyUpdatedAt)
- require.True(t, merged.Partial)
- require.Equal(t, []string{"weekly"}, merged.FailedWindows)
- require.Equal(t, []string{"monthly"}, previous.FailedWindows)
- }
- func floatPointer(value float64) *float64 {
- return &value
- }
|