1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
|
package controlplane
import (
"context"
"crypto/sha256"
"fmt"
"os"
"testing"
"time"
"aigw/internal/domain"
"github.com/jackc/pgx/v5/pgxpool"
)
func TestUsageFiltersAndDailyAggregationPostgres(t *testing.T) {
databaseURL := os.Getenv("AIGW_TEST_DATABASE_URL")
if databaseURL == "" {
t.Skip("AIGW_TEST_DATABASE_URL is not set")
}
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
if err := MigrateDatabase(ctx, databaseURL); err != nil {
t.Fatal(err)
}
db, err := pgxpool.New(ctx, databaseURL)
if err != nil {
t.Fatal(err)
}
t.Cleanup(db.Close)
store := &Store{db: db}
suffix := time.Now().UnixNano()
var tenantID, projectID, keyID, providerID, otherTenantID, otherProjectID, otherKeyID string
if err := db.QueryRow(ctx, `INSERT INTO tenants (slug,name) VALUES ($1,'Usage integration') RETURNING id::text`, fmt.Sprintf("usage-%d", suffix)).Scan(&tenantID); err != nil {
t.Fatal(err)
}
if err := db.QueryRow(ctx, `INSERT INTO projects (tenant_id,slug,name) VALUES ($1,'production','Production') RETURNING id::text`, tenantID).Scan(&projectID); err != nil {
t.Fatal(err)
}
keyHash := sha256.Sum256([]byte(fmt.Sprint(suffix)))
if err := db.QueryRow(ctx, `INSERT INTO api_keys (tenant_id,project_id,name,key_prefix,key_hash) VALUES ($1,$2,'Production key','sk-usage',$3) RETURNING id::text`, tenantID, projectID, keyHash[:]).Scan(&keyID); err != nil {
t.Fatal(err)
}
if err := db.QueryRow(ctx, `INSERT INTO providers (slug,name,protocol,wire_api,base_url,api_key_ciphertext)
VALUES ($1,$2,'openai','responses','https://usage.test',$3) RETURNING id::text`, fmt.Sprintf("usage-provider-%d", suffix), fmt.Sprintf("Usage provider %d", suffix), []byte("encrypted-test-value")).Scan(&providerID); err != nil {
t.Fatal(err)
}
if err := db.QueryRow(ctx, `INSERT INTO tenants (slug,name) VALUES ($1,'Other usage tenant') RETURNING id::text`, fmt.Sprintf("usage-other-%d", suffix)).Scan(&otherTenantID); err != nil {
t.Fatal(err)
}
if err := db.QueryRow(ctx, `INSERT INTO projects (tenant_id,slug,name) VALUES ($1,'production','Other production') RETURNING id::text`, otherTenantID).Scan(&otherProjectID); err != nil {
t.Fatal(err)
}
otherKeyHash := sha256.Sum256([]byte(fmt.Sprintf("other-%d", suffix)))
if err := db.QueryRow(ctx, `INSERT INTO api_keys (tenant_id,project_id,name,key_prefix,key_hash) VALUES ($1,$2,'Other key','sk-other',$3) RETURNING id::text`, otherTenantID, otherProjectID, otherKeyHash[:]).Scan(&otherKeyID); err != nil {
t.Fatal(err)
}
t.Cleanup(func() {
cleanupCtx := context.Background()
for _, statement := range []struct{ query, arg string }{
{`DELETE FROM usage_monthly_rollups WHERE tenant_id=$1`, tenantID},
{`DELETE FROM usage_monthly_rollups WHERE tenant_id=$1`, otherTenantID},
{`DELETE FROM usage_events WHERE tenant_id=$1`, tenantID},
{`DELETE FROM usage_events WHERE tenant_id=$1`, otherTenantID},
{`DELETE FROM providers WHERE id=$1`, providerID},
{`DELETE FROM api_keys WHERE id=$1`, keyID},
{`DELETE FROM api_keys WHERE id=$1`, otherKeyID},
{`DELETE FROM projects WHERE id=$1`, projectID},
{`DELETE FROM projects WHERE id=$1`, otherProjectID},
{`DELETE FROM tenants WHERE id=$1`, tenantID},
{`DELETE FROM tenants WHERE id=$1`, otherTenantID},
} {
if _, cleanupErr := db.Exec(cleanupCtx, statement.query, statement.arg); cleanupErr != nil {
t.Errorf("cleanup usage integration data: %v", cleanupErr)
}
}
})
started := time.Now().UTC().Add(-time.Hour).Truncate(time.Second)
events := []domain.UsageEvent{
{RequestID: fmt.Sprintf("req_usage_ok_%d", suffix), TenantID: tenantID, ProjectID: projectID, KeyID: keyID, PublicModel: "openai/test", ProviderID: providerID, Protocol: domain.ProtocolOpenAI, StatusCode: 200, Success: true, UsageReported: true, StartedAt: started, DurationMS: 120, Usage: domain.Usage{InputTokens: 10, OutputTokens: 4, TotalTokens: 14, CacheReadInputTokens: 5}},
{RequestID: fmt.Sprintf("req_usage_error_%d", suffix), TenantID: tenantID, ProjectID: projectID, KeyID: keyID, PublicModel: "openai/test", ProviderID: providerID, Protocol: domain.ProtocolOpenAIResponses, Stream: true, StatusCode: 502, Success: false, ErrorType: "provider_error", StartedAt: started.Add(time.Minute), DurationMS: 350},
}
for _, event := range events {
if err := store.RecordUsage(ctx, event); err != nil {
t.Fatal(err)
}
}
if _, err := db.Exec(ctx, `UPDATE usage_events SET charged_micros=125,cost_micros=125 WHERE request_id=$1`, events[0].RequestID); err != nil {
t.Fatal(err)
}
records, err := store.ListUsage(ctx, UsageQuery{TenantID: tenantID, KeyID: keyID, Status: "success", From: started.Add(-time.Minute), To: started.Add(time.Hour), Limit: 20})
if err != nil {
t.Fatal(err)
}
if len(records) != 1 || records[0].RequestID != events[0].RequestID || records[0].ProjectName != "Production" || records[0].KeyName != "Production key" || records[0].MeteringStatus != "reported" {
t.Fatalf("unexpected filtered usage: %+v", records)
}
streaming := true
providerSlug := fmt.Sprintf("usage-provider-%d", suffix)
records, err = store.ListUsage(ctx, UsageQuery{TenantID: tenantID, Provider: providerSlug, Protocol: string(domain.ProtocolOpenAIResponses), ErrorType: "provider_error", Stream: &streaming, From: started.Add(-time.Minute), To: started.Add(time.Hour), Limit: 20})
if err != nil {
t.Fatal(err)
}
if len(records) != 1 || records[0].RequestID != events[1].RequestID {
t.Fatalf("unexpected provider/protocol/transport usage filter: %+v", records)
}
filteredPoints, err := store.UsageDaily(ctx, UsageQuery{TenantID: tenantID, Provider: providerSlug, Protocol: string(domain.ProtocolOpenAIResponses), ErrorType: "provider_error", Stream: &streaming, From: started.Add(-time.Minute), To: started.Add(time.Hour)})
if err != nil {
t.Fatal(err)
}
if len(filteredPoints) != 1 || filteredPoints[0].RequestCount != 1 || filteredPoints[0].SuccessfulRequests != 0 {
t.Fatalf("unexpected filtered daily usage: %+v", filteredPoints)
}
points, err := store.UsageDaily(ctx, UsageQuery{TenantID: tenantID, From: started.Add(-time.Hour), To: started.Add(2 * time.Hour)})
if err != nil {
t.Fatal(err)
}
if len(points) != 1 || points[0].RequestCount != 2 || points[0].SuccessfulRequests != 1 || points[0].TotalTokens != 14 || points[0].ChargedMicros != 125 || points[0].P95DurationMS < 120 {
t.Fatalf("unexpected daily usage: %+v", points)
}
previous := domain.UsageEvent{RequestID: fmt.Sprintf("req_usage_previous_%d", suffix), TenantID: tenantID, ProjectID: projectID, KeyID: keyID, PublicModel: "openai/test", ProviderID: providerID, Protocol: domain.ProtocolOpenAI, StatusCode: 200, Success: true, UsageReported: true, StartedAt: started.Add(-5 * time.Minute), DurationMS: 80, Usage: domain.Usage{InputTokens: 3, OutputTokens: 2, TotalTokens: 5}}
missing := domain.UsageEvent{RequestID: fmt.Sprintf("req_usage_missing_%d", suffix), TenantID: tenantID, ProjectID: projectID, KeyID: keyID, PublicModel: "openai/test", ProviderID: providerID, Protocol: domain.ProtocolOpenAI, StatusCode: 200, Success: true, StartedAt: started.Add(2 * time.Minute), DurationMS: 200}
otherTenantEvent := domain.UsageEvent{RequestID: fmt.Sprintf("req_usage_other_%d", suffix), TenantID: otherTenantID, ProjectID: otherProjectID, KeyID: otherKeyID, PublicModel: "openai/test", ProviderID: providerID, Protocol: domain.ProtocolOpenAI, StatusCode: 200, Success: true, UsageReported: true, StartedAt: started.Add(3 * time.Minute), DurationMS: 900, Usage: domain.Usage{InputTokens: 1000, OutputTokens: 1000, TotalTokens: 2000}}
for _, event := range []domain.UsageEvent{previous, missing, otherTenantEvent} {
if err := store.RecordUsage(ctx, event); err != nil {
t.Fatal(err)
}
}
if _, err := db.Exec(ctx, `UPDATE usage_events SET charged_micros=25,cost_micros=25 WHERE request_id=$1`, previous.RequestID); err != nil {
t.Fatal(err)
}
analytics, err := store.UsageAnalytics(ctx, UsageQuery{TenantID: tenantID, From: started.Add(-time.Minute), To: started.Add(10 * time.Minute)})
if err != nil {
t.Fatal(err)
}
if len(analytics.Models) != 1 || analytics.Models[0].RequestCount != 3 || analytics.Models[0].SuccessfulRequests != 2 || analytics.Models[0].ProviderCount != 1 || analytics.Models[0].CacheReadInputTokens != 5 || analytics.Models[0].MissingUsageRequests != 1 || analytics.Models[0].ChargedMicros != 125 || analytics.Models[0].PreviousChargedMicros != 25 || analytics.Models[0].ChargeChangePercent == nil || *analytics.Models[0].ChargeChangePercent != 400 {
t.Fatalf("unexpected model analytics: %+v", analytics.Models)
}
if len(analytics.Providers) != 1 || analytics.Providers[0].ProviderID != providerID || analytics.Providers[0].WireAPI != "responses" || analytics.Providers[0].RequestCount != 3 || analytics.Providers[0].MissingUsageRequests != 1 || analytics.Providers[0].P95DurationMS < 200 {
t.Fatalf("unexpected provider analytics: %+v", analytics.Providers)
}
filteredAnalytics, err := store.UsageAnalytics(ctx, UsageQuery{TenantID: tenantID, Provider: providerSlug, Protocol: string(domain.ProtocolOpenAIResponses), ErrorType: "provider_error", Stream: &streaming, From: started.Add(-time.Minute), To: started.Add(10 * time.Minute)})
if err != nil {
t.Fatal(err)
}
if len(filteredAnalytics.Models) != 1 || filteredAnalytics.Models[0].RequestCount != 1 || filteredAnalytics.Models[0].ErrorCount != 1 || len(filteredAnalytics.Providers) != 1 || filteredAnalytics.Providers[0].RequestCount != 1 {
t.Fatalf("unexpected filtered analytics: %+v", filteredAnalytics)
}
}
|
