mirror of
https://github.com/CJackHwang/ds2api.git
synced 2026-05-08 18:35:35 +08:00
84 lines
2.5 KiB
Go
84 lines
2.5 KiB
Go
package openai
|
|
|
|
import (
|
|
"ds2api/internal/auth"
|
|
"net/http/httptest"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestIsVercelStreamPrepareRequest(t *testing.T) {
|
|
req := httptest.NewRequest("POST", "/v1/chat/completions?__stream_prepare=1", nil)
|
|
if !isVercelStreamPrepareRequest(req) {
|
|
t.Fatalf("expected prepare request to be detected")
|
|
}
|
|
|
|
req2 := httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
|
if isVercelStreamPrepareRequest(req2) {
|
|
t.Fatalf("expected non-prepare request")
|
|
}
|
|
}
|
|
|
|
func TestIsVercelStreamReleaseRequest(t *testing.T) {
|
|
req := httptest.NewRequest("POST", "/v1/chat/completions?__stream_release=1", nil)
|
|
if !isVercelStreamReleaseRequest(req) {
|
|
t.Fatalf("expected release request to be detected")
|
|
}
|
|
|
|
req2 := httptest.NewRequest("POST", "/v1/chat/completions", nil)
|
|
if isVercelStreamReleaseRequest(req2) {
|
|
t.Fatalf("expected non-release request")
|
|
}
|
|
}
|
|
|
|
func TestVercelInternalSecret(t *testing.T) {
|
|
t.Run("prefer explicit secret", func(t *testing.T) {
|
|
t.Setenv("DS2API_VERCEL_INTERNAL_SECRET", "stream-secret")
|
|
t.Setenv("DS2API_ADMIN_KEY", "admin-fallback")
|
|
if got := vercelInternalSecret(); got != "stream-secret" {
|
|
t.Fatalf("expected explicit secret, got %q", got)
|
|
}
|
|
})
|
|
|
|
t.Run("fallback to admin key", func(t *testing.T) {
|
|
t.Setenv("DS2API_VERCEL_INTERNAL_SECRET", "")
|
|
t.Setenv("DS2API_ADMIN_KEY", "admin-fallback")
|
|
if got := vercelInternalSecret(); got != "admin-fallback" {
|
|
t.Fatalf("expected admin key fallback, got %q", got)
|
|
}
|
|
})
|
|
|
|
t.Run("default admin when env missing", func(t *testing.T) {
|
|
t.Setenv("DS2API_VERCEL_INTERNAL_SECRET", "")
|
|
t.Setenv("DS2API_ADMIN_KEY", "")
|
|
if got := vercelInternalSecret(); got != "admin" {
|
|
t.Fatalf("expected default admin fallback, got %q", got)
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestStreamLeaseLifecycle(t *testing.T) {
|
|
h := &Handler{}
|
|
leaseID := h.holdStreamLease(&auth.RequestAuth{UseConfigToken: false})
|
|
if leaseID == "" {
|
|
t.Fatalf("expected non-empty lease id")
|
|
}
|
|
if ok := h.releaseStreamLease(leaseID); !ok {
|
|
t.Fatalf("expected lease release success")
|
|
}
|
|
if ok := h.releaseStreamLease(leaseID); ok {
|
|
t.Fatalf("expected duplicate release to fail")
|
|
}
|
|
}
|
|
|
|
func TestStreamLeaseTTL(t *testing.T) {
|
|
t.Setenv("DS2API_VERCEL_STREAM_LEASE_TTL_SECONDS", "120")
|
|
if got := streamLeaseTTL(); got != 120*time.Second {
|
|
t.Fatalf("expected ttl=120s, got %v", got)
|
|
}
|
|
t.Setenv("DS2API_VERCEL_STREAM_LEASE_TTL_SECONDS", "invalid")
|
|
if got := streamLeaseTTL(); got != 15*time.Minute {
|
|
t.Fatalf("expected default ttl on invalid value, got %v", got)
|
|
}
|
|
}
|