package test import ( "context" "fmt" "net/http" "net/http/httptest" "testing" "time" "github.com/router-for-me/CLIProxyAPI/v6/internal/config" runtimeexecutor "github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/executor" internalusage "github.com/router-for-me/CLIProxyAPI/v6/internal/usage" cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth" cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor" sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator") func TestGeminiExecutorRecordsSuccessfulZeroUsageInStatistics(t *testing.T) { model := fmt.Sprintf("gemini-2.5-flash-zero-usage-%d", time.Now().UnixNano()) source := fmt.Sprintf("zero-usage-%d@example.com", time.Now().UnixNano()) server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { wantPath := "/v1beta/models/" + model + ":generateContent" if r.URL.Path != wantPath { t.Fatalf("path = %q, want %q", r.URL.Path, wantPath) } w.Header().Set("Content-Type", "application/json") _, _ = w.Write([]byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":0,"candidatesTokenCount":0,"totalTokenCount":0}}`)) })) defer server.Close() executor := runtimeexecutor.NewGeminiExecutor(&config.Config{}) auth := &cliproxyauth.Auth{ Provider: "gemini", Attributes: map[string]string{ "api_key": "test-upstream-key", "base_url": server.URL, }, Metadata: map[string]any{ "email": source, }, } prevStatsEnabled := internalusage.StatisticsEnabled() internalusage.SetStatisticsEnabled(true) t.Cleanup(func() { internalusage.SetStatisticsEnabled(prevStatsEnabled) }) _, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{ Model: model, Payload: []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`), }, cliproxyexecutor.Options{ SourceFormat: sdktranslator.FormatGemini, OriginalRequest: []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`), }) if err != nil { t.Fatalf("Execute error: %v", err) } detail := waitForStatisticsDetail(t, "gemini", model, source) if detail.Failed { t.Fatalf("detail failed = true, want false") } if detail.Tokens.TotalTokens != 0 { t.Fatalf("total tokens = %d, want 0", detail.Tokens.TotalTokens) }} func waitForStatisticsDetail(t *testing.T, apiName, model, source string) internalusage.RequestDetail { t.Helper() deadline := time.Now().Add(2 * time.Second) for time.Now().Before(deadline) { snapshot := internalusage.GetRequestStatistics().Snapshot() apiSnapshot, ok := snapshot.APIs[apiName] if !ok { time.Sleep(10 * time.Millisecond) continue } modelSnapshot, ok := apiSnapshot.Models[model] if !ok { time.Sleep(10 * time.Millisecond) continue } for _, detail := range modelSnapshot.Details { if detail.Source == source { return detail } } time.Sleep(10 * time.Millisecond) } t.Fatalf("timed out waiting for statistics detail for api=%q model=%q source=%q", apiName, model, source) return internalusage.RequestDetail{}}
package test
import (
"context"
"fmt"
"net/http"
"net/http/httptest"
"testing"
"time"
"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
runtimeexecutor "github.com/router-for-me/CLIProxyAPI/v6/internal/runtime/executor"
internalusage "github.com/router-for-me/CLIProxyAPI/v6/internal/usage"
cliproxyauth "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/auth"
cliproxyexecutor "github.com/router-for-me/CLIProxyAPI/v6/sdk/cliproxy/executor"
sdktranslator "github.com/router-for-me/CLIProxyAPI/v6/sdk/translator"
)
func TestGeminiExecutorRecordsSuccessfulZeroUsageInStatistics(t *testing.T) {
model := fmt.Sprintf("gemini-2.5-flash-zero-usage-%d", time.Now().UnixNano())
source := fmt.Sprintf("zero-usage-%d@example.com", time.Now().UnixNano())
server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
wantPath := "/v1beta/models/" + model + ":generateContent"
if r.URL.Path != wantPath {
t.Fatalf("path = %q, want %q", r.URL.Path, wantPath)
}
w.Header().Set("Content-Type", "application/json")
_, _ = w.Write([]byte(`{"candidates":[{"content":{"role":"model","parts":[{"text":"ok"}]},"finishReason":"STOP"}],"usageMetadata":{"promptTokenCount":0,"candidatesTokenCount":0,"totalTokenCount":0}}`))
}))
defer server.Close()
executor := runtimeexecutor.NewGeminiExecutor(&config.Config{})
auth := &cliproxyauth.Auth{
Provider: "gemini",
Attributes: map[string]string{
"api_key": "test-upstream-key",
"base_url": server.URL,
},
Metadata: map[string]any{
"email": source,
},
}
prevStatsEnabled := internalusage.StatisticsEnabled()
internalusage.SetStatisticsEnabled(true)
t.Cleanup(func() {
internalusage.SetStatisticsEnabled(prevStatsEnabled)
})
_, err := executor.Execute(context.Background(), auth, cliproxyexecutor.Request{
Model: model,
Payload: []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`),
}, cliproxyexecutor.Options{
SourceFormat: sdktranslator.FormatGemini,
OriginalRequest: []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`),
})
if err != nil {
t.Fatalf("Execute error: %v", err)
}
detail := waitForStatisticsDetail(t, "gemini", model, source)
if detail.Failed {
t.Fatalf("detail failed = true, want false")
}
if detail.Tokens.TotalTokens != 0 {
t.Fatalf("total tokens = %d, want 0", detail.Tokens.TotalTokens)
}
}
func waitForStatisticsDetail(t *testing.T, apiName, model, source string) internalusage.RequestDetail {
t.Helper()
deadline := time.Now().Add(2 * time.Second)
for time.Now().Before(deadline) {
snapshot := internalusage.GetRequestStatistics().Snapshot()
apiSnapshot, ok := snapshot.APIs[apiName]
if !ok {
time.Sleep(10 * time.Millisecond)
continue
}
modelSnapshot, ok := apiSnapshot.Models[model]
if !ok {
time.Sleep(10 * time.Millisecond)
continue
}
for _, detail := range modelSnapshot.Details {
if detail.Source == source {
return detail
}
}
time.Sleep(10 * time.Millisecond)
}
t.Fatalf("timed out waiting for statistics detail for api=%q model=%q source=%q", apiName, model, source)
return internalusage.RequestDetail{}
}