add some scriptsmain

1 files changed, 83 insertions, 0 deletions

diff --git a/internal/httpapi/api_test.go b/internal/httpapi/api_test.go
index 66a344b..014b99b 100644
--- a/internal/httpapi/api_test.go
+++ b/internal/httpapi/api_test.go
@@ -3,6 +3,7 @@ package httpapi
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"encoding/json"
 	"io"
 	"log/slog"
@@ -14,6 +15,7 @@ import (
 	"time"
 
 	"aigw/internal/auth"
+	"aigw/internal/billing"
 	"aigw/internal/catalog"
 	"aigw/internal/config"
 	"aigw/internal/domain"
@@ -26,6 +28,20 @@ type captureUsageSink struct {
 	events chan domain.UsageEvent
 }
 
+type fakeBillingMeter struct {
+	authorizeErr error
+	settled      chan domain.UsageEvent
+}
+
+func (m *fakeBillingMeter) Authorize(context.Context, billing.Authorization) error {
+	return m.authorizeErr
+}
+
+func (m *fakeBillingMeter) Settle(_ context.Context, event domain.UsageEvent) error {
+	m.settled <- event
+	return nil
+}
+
 func (s *captureUsageSink) Publish(event domain.UsageEvent) {
 	s.events <- event
 }
@@ -178,7 +194,73 @@ func TestAnthropicHeadersAndPath(t *testing.T) {
 	}
 }
 
+func TestInsufficientBalanceRejectsBeforeCallingUpstream(t *testing.T) {
+	var calls atomic.Int64
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		calls.Add(1)
+		w.WriteHeader(http.StatusOK)
+	}))
+	defer upstream.Close()
+	meter := &fakeBillingMeter{authorizeErr: billing.ErrInsufficientBalance, settled: make(chan domain.UsageEvent, 1)}
+	gateway, sink := newTestGatewayWithBilling(t, []config.ProviderConfig{{
+		ID: "primary", Protocol: domain.ProtocolOpenAI, BaseURL: upstream.URL + "/v1", APIKey: "secret",
+	}}, []config.RouteConfig{{Provider: "primary", UpstreamModel: "model", Weight: 1}}, meter)
+	defer gateway.Close()
+
+	response := postOpenAI(t, gateway.URL, false)
+	defer response.Body.Close()
+	if response.StatusCode != http.StatusPaymentRequired {
+		t.Fatalf("status = %d, want 402", response.StatusCode)
+	}
+	if calls.Load() != 0 {
+		t.Fatalf("upstream calls = %d, want 0", calls.Load())
+	}
+	select {
+	case <-meter.settled:
+		t.Fatal("rejected request was settled")
+	default:
+	}
+	select {
+	case event := <-sink.events:
+		if event.StatusCode != http.StatusPaymentRequired || event.ErrorType != "insufficient_balance" {
+			t.Fatalf("unexpected rejection usage: %+v", event)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("rejected request did not emit usage")
+	}
+}
+
+func TestSuccessfulRequestIsSettled(t *testing.T) {
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		_, _ = io.WriteString(w, `{"choices":[],"usage":{"prompt_tokens":2,"completion_tokens":3,"total_tokens":5}}`)
+	}))
+	defer upstream.Close()
+	meter := &fakeBillingMeter{settled: make(chan domain.UsageEvent, 1)}
+	gateway, _ := newTestGatewayWithBilling(t, []config.ProviderConfig{{
+		ID: "primary", Protocol: domain.ProtocolOpenAI, BaseURL: upstream.URL + "/v1", APIKey: "secret",
+	}}, []config.RouteConfig{{Provider: "primary", UpstreamModel: "model", Weight: 1}}, meter)
+	defer gateway.Close()
+
+	response := postOpenAI(t, gateway.URL, false)
+	defer response.Body.Close()
+	if response.StatusCode != http.StatusOK {
+		t.Fatalf("status = %d, want 200", response.StatusCode)
+	}
+	select {
+	case event := <-meter.settled:
+		if event.Usage.TotalTokens != 5 {
+			t.Fatalf("settled usage = %+v", event.Usage)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("request was not settled")
+	}
+}
+
 func newTestGateway(t *testing.T, providers []config.ProviderConfig, routes []config.RouteConfig) (*httptest.Server, *captureUsageSink) {
+	return newTestGatewayWithBilling(t, providers, routes, nil)
+}
+
+func newTestGatewayWithBilling(t *testing.T, providers []config.ProviderConfig, routes []config.RouteConfig, meter billing.Meter) (*httptest.Server, *captureUsageSink) {
 	t.Helper()
 	authenticator, err := auth.NewStatic(`[{"key":"client-secret","key_id":"key-1","tenant_id":"tenant-1","project_id":"project-1","scopes":["inference"]}]`, false)
 	if err != nil {
@@ -201,6 +283,7 @@ func newTestGateway(t *testing.T, providers []config.ProviderConfig, routes []co
 		Router:        routing.New(modelCatalog),
 		Forwarder:     provider.New(cfg.UpstreamHTTP, metrics),
 		UsageSink:     sink,
+		BillingMeter:  meter,
 		Metrics:       metrics,
 		Logger:        slog.New(slog.NewTextHandler(io.Discard, nil)),
 		MaxBodyBytes:  1 << 20,