internal/api/middleware/response_cache_test.go

package middleware

import (
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"

	"github.com/gin-gonic/gin"
	"github.com/router-for-me/CLIProxyAPI/v6/internal/cache"
	"github.com/router-for-me/CLIProxyAPI/v6/internal/config"
)

func testResponseCacheConfig() config.ResponseCacheConfig {
	return config.ResponseCacheConfig{
		Enabled:           true,
		TTLSeconds:        300,
		MaxMemoryEntries:  100,
		MaxResponseSizeKB: 512,
		ExcludeStreaming:  true,
		RedisKeyPrefix:    "test",
	}
}

func setupRouterWithCache(t *testing.T) *gin.Engine {
	t.Helper()
	gin.SetMode(gin.TestMode)
	r := gin.New()
	cfg := testResponseCacheConfig()
	rc := cache.NewResponseCache(cfg, nil)
	r.Use(ResponseCacheMiddleware(rc, cfg))
	return r
}

func TestResponseCacheMiddleware_CachesAndHitsNonStreaming(t *testing.T) {
	r := setupRouterWithCache(t)
	calls := 0
	r.POST("/v1/chat/completions", func(c *gin.Context) {
		calls++
		c.JSON(http.StatusOK, gin.H{"ok": true, "n": calls})
	})

	body := `{"model":"gpt-5","messages":[{"role":"user","content":"hi"}]}`

	req1 := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body))
	req1.Header.Set("Content-Type", "application/json")
	w1 := httptest.NewRecorder()
	r.ServeHTTP(w1, req1)
	if w1.Code != http.StatusOK {
		t.Fatalf("first request status: got %d", w1.Code)
	}
	if got := w1.Header().Get("X-Cache"); got != "" {
		t.Fatalf("first request should not be cache hit, got X-Cache=%q", got)
	}

	req2 := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body))
	req2.Header.Set("Content-Type", "application/json")
	w2 := httptest.NewRecorder()
	r.ServeHTTP(w2, req2)
	if w2.Code != http.StatusOK {
		t.Fatalf("second request status: got %d", w2.Code)
	}
	if got := w2.Header().Get("X-Cache"); got != "HIT" {
		t.Fatalf("second request should be cache hit, got X-Cache=%q", got)
	}
	if calls != 1 {
		t.Fatalf("handler should be called once, got %d", calls)
	}
}

func TestResponseCacheMiddleware_DoesNotCacheStreaming(t *testing.T) {
	r := setupRouterWithCache(t)
	calls := 0
	r.POST("/v1/chat/completions", func(c *gin.Context) {
		calls++
		c.JSON(http.StatusOK, gin.H{"ok": true, "n": calls})
	})

	body := `{"model":"gpt-5","stream":true,"messages":[{"role":"user","content":"hi"}]}`

	req1 := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body))
	req1.Header.Set("Content-Type", "application/json")
	w1 := httptest.NewRecorder()
	r.ServeHTTP(w1, req1)

	req2 := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body))
	req2.Header.Set("Content-Type", "application/json")
	w2 := httptest.NewRecorder()
	r.ServeHTTP(w2, req2)

	if calls != 2 {
		t.Fatalf("streaming requests should not be cached; handler calls=%d", calls)
	}
	if got := w2.Header().Get("X-Cache"); got != "" {
		t.Fatalf("streaming request should not have X-Cache hit, got %q", got)
	}
}

func TestResponseCacheMiddleware_SkipsNonJSONContentType(t *testing.T) {
	r := setupRouterWithCache(t)
	calls := 0
	r.POST("/v1/chat/completions", func(c *gin.Context) {
		calls++
		c.JSON(http.StatusOK, gin.H{"ok": true, "n": calls})
	})

	body := "model=gpt-5&q=hello"
	req1 := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body))
	req1.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	w1 := httptest.NewRecorder()
	r.ServeHTTP(w1, req1)

	req2 := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(body))
	req2.Header.Set("Content-Type", "application/x-www-form-urlencoded")
	w2 := httptest.NewRecorder()
	r.ServeHTTP(w2, req2)

	if calls != 2 {
		t.Fatalf("non-json requests should bypass cache; handler calls=%d", calls)
	}
	if got := w2.Header().Get("X-Cache"); got != "" {
		t.Fatalf("non-json request should not have X-Cache hit, got %q", got)
	}
}