sub2api/backend/internal/handler/gemini_cached_contents_handler_test.go

//go:build unit

package handler

import (
	"context"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/Wei-Shaw/sub2api/internal/config"
	middleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
	"github.com/Wei-Shaw/sub2api/internal/service"
	"github.com/gin-gonic/gin"
	"github.com/stretchr/testify/require"
	"go.uber.org/zap"
)

type geminiCachedContentHandlerRepoStub struct {
	records []*service.GeminiCachedContent
}

func (r *geminiCachedContentHandlerRepoStub) Create(_ context.Context, record *service.GeminiCachedContent) error {
	record.ID = int64(len(r.records) + 1)
	r.records = append(r.records, record)
	return nil
}

func (r *geminiCachedContentHandlerRepoStub) GetForOwner(_ context.Context, apiKeyID int64, publicID string) (*service.GeminiCachedContent, error) {
	for _, record := range r.records {
		if record.PublicID == publicID && record.APIKeyID == apiKeyID {
			return record, nil
		}
	}
	return nil, service.ErrGeminiCachedContentNotFound
}

func (r *geminiCachedContentHandlerRepoStub) ListForOwner(_ context.Context, apiKeyID, groupID int64, activeAt time.Time, beforeID int64, limit int) ([]*service.GeminiCachedContent, error) {
	var out []*service.GeminiCachedContent
	for i := len(r.records) - 1; i >= 0; i-- {
		record := r.records[i]
		if record.APIKeyID == apiKeyID && record.GroupID == groupID && record.ExpireTime.After(activeAt) && (beforeID <= 0 || record.ID < beforeID) {
			out = append(out, record)
		}
		if len(out) == limit {
			break
		}
	}
	return out, nil
}

func (r *geminiCachedContentHandlerRepoStub) UpdateExpireTime(context.Context, int64, time.Time) (time.Time, error) {
	return time.Time{}, nil
}

func (r *geminiCachedContentHandlerRepoStub) SoftDelete(context.Context, int64) error {
	return nil
}

func newGeminiCachedContentHandlerForTest(repo *geminiCachedContentHandlerRepoStub) *GatewayHandler {
	return &GatewayHandler{geminiCachedContentService: service.NewGeminiCachedContentService(repo, nil, &config.Config{})}
}

func geminiCachedContentTestAPIKey(platform string) *service.APIKey {
	groupID := int64(42)
	return &service.APIKey{ID: 7, UserID: 3, GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: platform}}
}

func newGeminiCachedContentTestContext(method, target string, apiKey *service.APIKey) (*gin.Context, *httptest.ResponseRecorder) {
	gin.SetMode(gin.TestMode)
	rec := httptest.NewRecorder()
	c, _ := gin.CreateTestContext(rec)
	c.Request = httptest.NewRequest(method, target, nil)
	c.Set(string(middleware.ContextKeyAPIKey), apiKey)
	c.Set(string(middleware.ContextKeyUser), middleware.AuthSubject{UserID: apiKey.UserID, Concurrency: 10})
	return c, rec
}

func seedGeminiCachedContent(repo *geminiCachedContentHandlerRepoStub, publicID string, apiKeyID int64, expire time.Time) *service.GeminiCachedContent {
	record := &service.GeminiCachedContent{
		PublicID:        publicID,
		UserID:          3,
		APIKeyID:        apiKeyID,
		GroupID:         42,
		AccountID:       900,
		UpstreamName:    "cachedContents/upstream" + publicID,
		Model:           "gemini-3.8-flash",
		UpstreamModel:   "models/gemini-3.8-flash",
		TotalTokenCount: 6273,
		ExpireTime:      expire,
		CreatedAt:       time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC),
		UpdatedAt:       time.Date(2026, 10, 9, 8, 0, 0, 0, time.UTC),
	}
	_ = repo.Create(context.Background(), record)
	return record
}

func TestResolveGeminiCachedContentBinding(t *testing.T) {
	repo := &geminiCachedContentHandlerRepoStub{}
	seedGeminiCachedContent(repo, "abc", 7, time.Now().Add(time.Hour))
	seedGeminiCachedContent(repo, "foreign", 8, time.Now().Add(time.Hour))
	h := newGeminiCachedContentHandlerForTest(repo)
	apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini)

	t.Run("request without cache passes through", func(t *testing.T) {
		c, _ := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
		body := []byte(`{"contents":[]}`)
		record, out, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), apiKey, "gemini-3.8-flash", body)
		require.True(t, ok)
		require.Nil(t, record)
		require.Equal(t, body, out)
	})

	t.Run("owned cache rewrites reference and binds account", func(t *testing.T) {
		c, _ := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
		record, out, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), apiKey, "gemini-3.8-flash",
			[]byte(`{"cachedContent":"cachedContents/abc","contents":[]}`))
		require.True(t, ok)
		require.Equal(t, int64(900), record.AccountID)
		require.JSONEq(t, `{"cachedContent":"cachedContents/upstreamabc","contents":[]}`, string(out))
	})

	t.Run("cache owned by another key looks missing", func(t *testing.T) {
		c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
		_, _, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), apiKey, "gemini-3.8-flash",
			[]byte(`{"cachedContent":"cachedContents/foreign","contents":[]}`))
		require.False(t, ok)
		require.Equal(t, http.StatusForbidden, rec.Code)
		require.Contains(t, rec.Body.String(), service.GeminiCachedContentNotFoundMessage)
	})

	t.Run("model must match the cache literally", func(t *testing.T) {
		c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-flash-latest:generateContent", apiKey)
		_, _, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), apiKey, "gemini-flash-latest",
			[]byte(`{"cachedContent":"cachedContents/abc","contents":[]}`))
		require.False(t, ok)
		require.Equal(t, http.StatusBadRequest, rec.Code)
		require.Contains(t, rec.Body.String(), "has to be the same")
	})

	t.Run("composite groups are rejected", func(t *testing.T) {
		composite := geminiCachedContentTestAPIKey(service.PlatformComposite)
		c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", composite)
		_, _, ok := h.resolveGeminiCachedContentBinding(c, zap.NewNop(), composite, "gemini-3.8-flash",
			[]byte(`{"cachedContent":"cachedContents/abc","contents":[]}`))
		require.False(t, ok)
		require.Equal(t, http.StatusBadRequest, rec.Code)
	})
}

func TestGeminiCachedContentsListAndGet(t *testing.T) {
	repo := &geminiCachedContentHandlerRepoStub{}
	seedGeminiCachedContent(repo, "first", 7, time.Now().Add(time.Hour))
	seedGeminiCachedContent(repo, "second", 7, time.Now().Add(time.Hour))
	seedGeminiCachedContent(repo, "expired", 7, time.Now().Add(-time.Minute))
	seedGeminiCachedContent(repo, "foreign", 8, time.Now().Add(time.Hour))
	h := newGeminiCachedContentHandlerForTest(repo)
	apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini)

	c, rec := newGeminiCachedContentTestContext(http.MethodGet, "/v1beta/cachedContents?pageSize=1", apiKey)
	h.GeminiCachedContentsList(c)
	require.Equal(t, http.StatusOK, rec.Code)
	var page struct {
		CachedContents []map[string]any `json:"cachedContents"`
		NextPageToken  string           `json:"nextPageToken"`
	}
	require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &page))
	require.Len(t, page.CachedContents, 1)
	require.Equal(t, "cachedContents/second", page.CachedContents[0]["name"])
	require.NotEmpty(t, page.NextPageToken)

	c, rec = newGeminiCachedContentTestContext(http.MethodGet, "/v1beta/cachedContents?pageSize=1&pageToken="+page.NextPageToken, apiKey)
	h.GeminiCachedContentsList(c)
	require.Equal(t, http.StatusOK, rec.Code)
	page.NextPageToken = ""
	require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &page))
	require.Len(t, page.CachedContents, 1)
	require.Equal(t, "cachedContents/first", page.CachedContents[0]["name"])
	require.Empty(t, page.NextPageToken, "过期与其他 Key 的缓存不出现在列表里")

	c, rec = newGeminiCachedContentTestContext(http.MethodGet, "/v1beta/cachedContents/first", apiKey)
	c.Params = gin.Params{{Key: "cacheID", Value: "first"}}
	h.GeminiCachedContentsGet(c)
	require.Equal(t, http.StatusOK, rec.Code)
	var view map[string]any
	require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &view))
	require.Equal(t, "cachedContents/first", view["name"])
	require.Equal(t, "models/gemini-3.8-flash", view["model"])
	require.NotContains(t, rec.Body.String(), "upstreamfirst", "上游资源名不暴露给客户端")

	c, rec = newGeminiCachedContentTestContext(http.MethodGet, "/v1beta/cachedContents/foreign", apiKey)
	c.Params = gin.Params{{Key: "cacheID", Value: "foreign"}}
	h.GeminiCachedContentsGet(c)
	require.Equal(t, http.StatusForbidden, rec.Code)
	require.True(t, strings.Contains(rec.Body.String(), service.GeminiCachedContentNotFoundMessage))
}

func TestHandleGeminiCachedContentFailoverExhausted(t *testing.T) {
	h := &GatewayHandler{}
	apiKey := geminiCachedContentTestAPIKey(service.PlatformGemini)
	unavailable := `{"error":{"code":503,"message":"The account holding this cached content is temporarily unavailable","status":"UNAVAILABLE"}}`

	c, rec := newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
	h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusServiceUnavailable, ResponseBody: []byte(unavailable)})
	require.Equal(t, http.StatusServiceUnavailable, rec.Code, "绑定缓存的请求无账号可换，上游 503 原样返回")
	require.JSONEq(t, unavailable, rec.Body.String())

	c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
	h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusTooManyRequests, ResponseBody: []byte(`{"error":{"code":429,"message":"Too many pending requests, please retry later","status":"RESOURCE_EXHAUSTED"}}`)})
	require.Equal(t, http.StatusTooManyRequests, rec.Code)
	require.Contains(t, rec.Body.String(), "Too many pending requests")

	c, rec = newGeminiCachedContentTestContext(http.MethodPost, "/v1beta/models/gemini-3.8-flash:generateContent", apiKey)
	h.handleGeminiCachedContentFailoverExhausted(c, &service.UpstreamFailoverError{StatusCode: http.StatusUnauthorized, ResponseBody: []byte(`{"error":{"code":401,"message":"bad key"}}`)})
	require.Equal(t, http.StatusBadGateway, rec.Code, "账号凭据类错误仍按网关映射，不透传给客户端")
	require.NotContains(t, rec.Body.String(), "bad key")
}