package service import ( "testing" "github.com/Wei-Shaw/sub2api/internal/pkg/claude" "github.com/stretchr/testify/require" ) func TestMergeAnthropicBeta(t *testing.T) { got := mergeAnthropicBeta( []string{"oauth-2025-04-20", "interleaved-thinking-2025-05-14"}, "foo, oauth-2025-04-20,bar, foo", ) require.Equal(t, "oauth-2025-04-20,interleaved-thinking-2025-05-14,foo,bar", got) } func TestMergeAnthropicBeta_EmptyIncoming(t *testing.T) { got := mergeAnthropicBeta( []string{"oauth-2025-04-20", "interleaved-thinking-2025-05-14"}, "", ) require.Equal(t, "oauth-2025-04-20,interleaved-thinking-2025-05-14", got) } func TestStripBetaTokens(t *testing.T) { tests := []struct { name string header string tokens []string want string }{ { name: "single token in middle", header: "oauth-2025-04-20,context-1m-2025-08-07,interleaved-thinking-2025-05-14", tokens: []string{"context-1m-2025-08-07"}, want: "oauth-2025-04-20,interleaved-thinking-2025-05-14", }, { name: "single token at start", header: "context-1m-2025-08-07,oauth-2025-04-20,interleaved-thinking-2025-05-14", tokens: []string{"context-1m-2025-08-07"}, want: "oauth-2025-04-20,interleaved-thinking-2025-05-14", }, { name: "single token at end", header: "oauth-2025-04-20,interleaved-thinking-2025-05-14,context-1m-2025-08-07", tokens: []string{"context-1m-2025-08-07"}, want: "oauth-2025-04-20,interleaved-thinking-2025-05-14", }, { name: "token not present", header: "oauth-2025-04-20,interleaved-thinking-2025-05-14", tokens: []string{"context-1m-2025-08-07"}, want: "oauth-2025-04-20,interleaved-thinking-2025-05-14", }, { name: "empty header", header: "", tokens: []string{"context-1m-2025-08-07"}, want: "", }, { name: "with spaces", header: "oauth-2025-04-20, context-1m-2025-08-07 , interleaved-thinking-2025-05-14", tokens: []string{"context-1m-2025-08-07"}, want: "oauth-2025-04-20,interleaved-thinking-2025-05-14", }, { name: "only token", header: "context-1m-2025-08-07", tokens: []string{"context-1m-2025-08-07"}, want: "", }, { name: "nil tokens", header: "oauth-2025-04-20,interleaved-thinking-2025-05-14", tokens: nil, want: "oauth-2025-04-20,interleaved-thinking-2025-05-14", }, { name: "multiple tokens removed", header: "oauth-2025-04-20,context-1m-2025-08-07,interleaved-thinking-2025-05-14,fast-mode-2026-02-01", tokens: []string{"context-1m-2025-08-07", "fast-mode-2026-02-01"}, want: "oauth-2025-04-20,interleaved-thinking-2025-05-14", }, { name: "DroppedBetas removes both context-1m and fast-mode", header: "oauth-2025-04-20,context-1m-2025-08-07,fast-mode-2026-02-01,interleaved-thinking-2025-05-14", tokens: claude.DroppedBetas, want: "oauth-2025-04-20,interleaved-thinking-2025-05-14", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got := stripBetaTokens(tt.header, tt.tokens) require.Equal(t, tt.want, got) }) } } func TestMergeAnthropicBetaDropping_Context1M(t *testing.T) { required := []string{"oauth-2025-04-20", "interleaved-thinking-2025-05-14"} incoming := "context-1m-2025-08-07,foo-beta,oauth-2025-04-20" drop := map[string]struct{}{"context-1m-2025-08-07": {}} got := mergeAnthropicBetaDropping(required, incoming, drop) require.Equal(t, "oauth-2025-04-20,interleaved-thinking-2025-05-14,foo-beta", got) require.NotContains(t, got, "context-1m-2025-08-07") } func TestMergeAnthropicBetaDropping_DroppedBetas(t *testing.T) { required := []string{"oauth-2025-04-20", "interleaved-thinking-2025-05-14"} incoming := "context-1m-2025-08-07,fast-mode-2026-02-01,foo-beta,oauth-2025-04-20" drop := droppedBetaSet() got := mergeAnthropicBetaDropping(required, incoming, drop) require.Equal(t, "oauth-2025-04-20,interleaved-thinking-2025-05-14,foo-beta", got) require.NotContains(t, got, "context-1m-2025-08-07") require.NotContains(t, got, "fast-mode-2026-02-01") } func TestDroppedBetaSet(t *testing.T) { // Base set contains DroppedBetas base := droppedBetaSet() require.Contains(t, base, claude.BetaContext1M) require.Contains(t, base, claude.BetaFastMode) require.Len(t, base, len(claude.DroppedBetas)) // With extra tokens extended := droppedBetaSet(claude.BetaClaudeCode) require.Contains(t, extended, claude.BetaContext1M) require.Contains(t, extended, claude.BetaFastMode) require.Contains(t, extended, claude.BetaClaudeCode) require.Len(t, extended, len(claude.DroppedBetas)+1) } func TestBuildBetaTokenSet(t *testing.T) { got := buildBetaTokenSet([]string{"foo", "", "bar", "foo"}) require.Len(t, got, 2) require.Contains(t, got, "foo") require.Contains(t, got, "bar") require.NotContains(t, got, "") empty := buildBetaTokenSet(nil) require.Empty(t, empty) } func TestStripBetaTokensWithSet_EmptyDropSet(t *testing.T) { header := "oauth-2025-04-20,interleaved-thinking-2025-05-14" got := stripBetaTokensWithSet(header, map[string]struct{}{}) require.Equal(t, header, got) } func TestIsCountTokensUnsupported404(t *testing.T) { tests := []struct { name string statusCode int body string want bool }{ { name: "exact endpoint not found", statusCode: 404, body: `{"error":{"message":"Not found: /v1/messages/count_tokens","type":"not_found_error"}}`, want: true, }, { name: "contains count_tokens and not found", statusCode: 404, body: `{"error":{"message":"count_tokens route not found","type":"not_found_error"}}`, want: true, }, { name: "generic 404", statusCode: 404, body: `{"error":{"message":"resource not found","type":"not_found_error"}}`, want: false, }, { name: "404 with empty error message", statusCode: 404, body: `{"error":{"message":"","type":"not_found_error"}}`, want: false, }, { name: "non-404 status", statusCode: 400, body: `{"error":{"message":"Not found: /v1/messages/count_tokens","type":"invalid_request_error"}}`, want: false, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { got := isCountTokensUnsupported404(tt.statusCode, []byte(tt.body)) require.Equal(t, tt.want, got) }) } }