gateway_handler_intercept_test.go 1.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465
  1. package handler
  2. import (
  3. "encoding/json"
  4. "net/http"
  5. "net/http/httptest"
  6. "strings"
  7. "testing"
  8. "github.com/gin-gonic/gin"
  9. "github.com/stretchr/testify/require"
  10. )
  11. func TestDetectInterceptType_MaxTokensOneHaikuRequiresClaudeCodeClient(t *testing.T) {
  12. body := []byte(`{"messages":[{"role":"user","content":[{"type":"text","text":"hello"}]}]}`)
  13. notClaudeCode := detectInterceptType(body, "claude-haiku-4-5", 1, false)
  14. require.Equal(t, InterceptTypeNone, notClaudeCode)
  15. isClaudeCode := detectInterceptType(body, "claude-haiku-4-5", 1, true)
  16. require.Equal(t, InterceptTypeMaxTokensOneHaiku, isClaudeCode)
  17. }
  18. func TestDetectInterceptType_SuggestionModeUnaffected(t *testing.T) {
  19. body := []byte(`{
  20. "messages":[{
  21. "role":"user",
  22. "content":[{"type":"text","text":"[SUGGESTION MODE:foo]"}]
  23. }],
  24. "system":[]
  25. }`)
  26. got := detectInterceptType(body, "claude-sonnet-4-5", 256, false)
  27. require.Equal(t, InterceptTypeSuggestionMode, got)
  28. }
  29. func TestSendMockInterceptResponse_MaxTokensOneHaiku(t *testing.T) {
  30. gin.SetMode(gin.TestMode)
  31. rec := httptest.NewRecorder()
  32. ctx, _ := gin.CreateTestContext(rec)
  33. sendMockInterceptResponse(ctx, "claude-haiku-4-5", InterceptTypeMaxTokensOneHaiku)
  34. require.Equal(t, http.StatusOK, rec.Code)
  35. var response map[string]any
  36. require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &response))
  37. require.Equal(t, "max_tokens", response["stop_reason"])
  38. id, ok := response["id"].(string)
  39. require.True(t, ok)
  40. require.True(t, strings.HasPrefix(id, "msg_01"))
  41. content, ok := response["content"].([]any)
  42. require.True(t, ok)
  43. require.NotEmpty(t, content)
  44. firstBlock, ok := content[0].(map[string]any)
  45. require.True(t, ok)
  46. require.Equal(t, "#", firstBlock["text"])
  47. usage, ok := response["usage"].(map[string]any)
  48. require.True(t, ok)
  49. require.Equal(t, float64(1), usage["output_tokens"])
  50. }