59 lines
1.8 KiB
Go
59 lines
1.8 KiB
Go
package service
|
|
|
|
import (
|
|
"context"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
|
"github.com/stretchr/testify/require"
|
|
)
|
|
|
|
func TestClaudeCodeValidator_ProbeBypass(t *testing.T) {
|
|
validator := NewClaudeCodeValidator()
|
|
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
|
|
req.Header.Set("User-Agent", "claude-cli/1.2.3 (darwin; arm64)")
|
|
req = req.WithContext(context.WithValue(req.Context(), ctxkey.IsMaxTokensOneHaikuRequest, true))
|
|
|
|
ok := validator.Validate(req, map[string]any{
|
|
"model": "claude-haiku-4-5",
|
|
"max_tokens": 1,
|
|
})
|
|
require.True(t, ok)
|
|
}
|
|
|
|
func TestClaudeCodeValidator_ProbeBypassRequiresUA(t *testing.T) {
|
|
validator := NewClaudeCodeValidator()
|
|
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
|
|
req.Header.Set("User-Agent", "curl/8.0.0")
|
|
req = req.WithContext(context.WithValue(req.Context(), ctxkey.IsMaxTokensOneHaikuRequest, true))
|
|
|
|
ok := validator.Validate(req, map[string]any{
|
|
"model": "claude-haiku-4-5",
|
|
"max_tokens": 1,
|
|
})
|
|
require.False(t, ok)
|
|
}
|
|
|
|
func TestClaudeCodeValidator_MessagesWithoutProbeStillNeedStrictValidation(t *testing.T) {
|
|
validator := NewClaudeCodeValidator()
|
|
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/messages", nil)
|
|
req.Header.Set("User-Agent", "claude-cli/1.2.3 (darwin; arm64)")
|
|
|
|
ok := validator.Validate(req, map[string]any{
|
|
"model": "claude-haiku-4-5",
|
|
"max_tokens": 1,
|
|
})
|
|
require.False(t, ok)
|
|
}
|
|
|
|
func TestClaudeCodeValidator_NonMessagesPathUAOnly(t *testing.T) {
|
|
validator := NewClaudeCodeValidator()
|
|
req := httptest.NewRequest(http.MethodPost, "http://example.com/v1/models", nil)
|
|
req.Header.Set("User-Agent", "claude-cli/1.2.3 (darwin; arm64)")
|
|
|
|
ok := validator.Validate(req, nil)
|
|
require.True(t, ok)
|
|
}
|