Sub2API v1.0 - AI API 网关（二开初始版本，基于上游 Wei-Shaw/sub2api）

backend/internal/service/account_usage_service_spark_shadow_test.go

@@ -0,0 +1,151 @@
+package service
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+)
+
+// sparkShadowUsageTestRepo is a minimal AccountRepository stub for spark shadow
+// usage tests.  GetByID serves both shadow and parent accounts from a map;
+// UpdateExtra records the persisted updates for assertion.
+type sparkShadowUsageTestRepo struct {
+	AccountRepository
+	accounts      map[int64]*Account
+	updateExtraCh chan map[string]any
+}
+
+func (r *sparkShadowUsageTestRepo) GetByID(_ context.Context, id int64) (*Account, error) {
+	if acc, ok := r.accounts[id]; ok {
+		return acc, nil
+	}
+	return nil, fmt.Errorf("account %d not found", id)
+}
+
+func (r *sparkShadowUsageTestRepo) UpdateExtra(_ context.Context, _ int64, updates map[string]any) error {
+	if r.updateExtraCh != nil {
+		copied := make(map[string]any, len(updates))
+		for k, v := range updates {
+			copied[k] = v
+		}
+		r.updateExtraCh <- copied
+	}
+	return nil
+}
+
+// TestGetOpenAIUsage_SparkShadow_WritesExtraAndReturnsNonEmptyWindows covers
+// two assertions required by Task 3.2:
+//
+// A) After getOpenAIUsage on a spark shadow account the shadow row's
+// Extra["codex_5h_used_percent"] is persisted, and the upstream call carried
+// the PARENT account's chatgpt-account-id (not the shadow's empty one).
+//
+// B) (P1-b regression guard) The UsageInfo RETURNED by the same call has
+// non-nil FiveHour AND SevenDay windows — proving that the rebuild happened
+// and not just the DB write.
+func TestGetOpenAIUsage_SparkShadow_WritesExtraAndReturnsNonEmptyWindows(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+
+	pid := int64(100)
+	shadow := &Account{
+		ID:              200,
+		ParentAccountID: &pid,
+		Platform:        PlatformOpenAI,
+		Type:            AccountTypeOAuth,
+		Status:          StatusActive,
+		QuotaDimension:  QuotaDimensionSpark,
+	}
+	parent := &Account{
+		ID:       100,
+		Platform: PlatformOpenAI,
+		Type:     AccountTypeOAuth,
+		Status:   StatusActive,
+		Credentials: map[string]any{
+			"chatgpt_account_id": "org-spark-parent",
+		},
+	}
+
+	// Repo shared by both the OpenAIQuotaService (needs shadow+parent for resolve)
+	// and the AccountUsageService (needs UpdateExtra for persist).
+	updateExtraCh := make(chan map[string]any, 1)
+	repo := &sparkShadowUsageTestRepo{
+		accounts:      map[int64]*Account{200: shadow, 100: parent},
+		updateExtraCh: updateExtraCh,
+	}
+
+	// Token cache: return a fake token for the parent account key.
+	tokenCache := &stubQuotaTokenCache{tokens: map[string]string{
+		OpenAITokenCacheKey(parent): "fake-access-token",
+	}}
+	tokenProvider := NewOpenAITokenProvider(repo, tokenCache, nil)
+
+	// httptest server: records the chatgpt-account-id header and returns a
+	// synthetic OpenAIQuotaUsage with codex_bengalfox 5h+7d windows.
+	var capturedAccountID string
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		capturedAccountID = r.Header.Get("chatgpt-account-id")
+		w.Header().Set("content-type", "application/json")
+		resp := OpenAIQuotaUsage{
+			AdditionalRateLimits: []OpenAIAdditionalRateLimit{
+				{
+					MeteredFeature: "codex_bengalfox",
+					RateLimit: &OpenAIRateLimit{
+						// Primary window → 5h (18000 s = 300 min)
+						PrimaryWindow: &OpenAIRateLimitWindow{
+							UsedPercent:        42.5,
+							ResetAfterSeconds:  3600,
+							LimitWindowSeconds: 18000,
+						},
+						// Secondary window → 7d (604800 s = 10080 min)
+						SecondaryWindow: &OpenAIRateLimitWindow{
+							UsedPercent:        10.0,
+							ResetAfterSeconds:  86400,
+							LimitWindowSeconds: 604800,
+						},
+					},
+				},
+			},
+		}
+		_ = json.NewEncoder(w).Encode(resp)
+	}))
+	defer srv.Close()
+
+	quotaService := NewOpenAIQuotaService(repo, nil, tokenProvider, newQuotaRedirectingFactory(srv))
+	svc := &AccountUsageService{
+		accountRepo:        repo,
+		openAIQuotaService: quotaService,
+	}
+
+	usage, err := svc.getOpenAIUsage(ctx, shadow, true /*force*/)
+	require.NoError(t, err)
+
+	// Assertion A-1: upstream received the PARENT's chatgpt-account-id.
+	require.Equal(t, "org-spark-parent", capturedAccountID,
+		"QueryUsage must use parent's chatgpt-account-id for spark shadow accounts")
+
+	// Assertion A-2: shadow Extra was persisted with codex_5h_used_percent.
+	select {
+	case updates := <-updateExtraCh:
+		require.Contains(t, updates, "codex_5h_used_percent",
+			"persisted extra must contain codex_5h_used_percent")
+		require.InDelta(t, 42.5, updates["codex_5h_used_percent"], 0.01,
+			"codex_5h_used_percent must match the upstream value")
+	case <-time.After(2 * time.Second):
+		t.Fatal("UpdateExtra was not called within timeout — spark shadow persist did not happen")
+	}
+
+	// Assertion B (P1-b regression guard): returned UsageInfo must have
+	// non-nil windows. This FAILS if the code only writes Extra without
+	// rebuilding the returned UsageInfo.
+	require.NotNil(t, usage.FiveHour,
+		"returned UsageInfo.FiveHour must be non-nil (rebuild from merged Extra must happen)")
+	require.NotNil(t, usage.SevenDay,
+		"returned UsageInfo.SevenDay must be non-nil (rebuild from merged Extra must happen)")
+}