req_client_pool_test.go 4.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142
  1. package repository
  2. import (
  3. "context"
  4. "net/http"
  5. "net/http/httptest"
  6. "reflect"
  7. "strings"
  8. "sync"
  9. "testing"
  10. "time"
  11. "unsafe"
  12. "github.com/Wei-Shaw/sub2api/internal/pkg/servertiming"
  13. "github.com/imroc/req/v3"
  14. "github.com/stretchr/testify/require"
  15. )
  16. func forceHTTPVersion(t *testing.T, client *req.Client) string {
  17. t.Helper()
  18. transport := client.GetTransport()
  19. field := reflect.ValueOf(transport).Elem().FieldByName("forceHttpVersion")
  20. require.True(t, field.IsValid(), "forceHttpVersion field not found")
  21. require.True(t, field.CanAddr(), "forceHttpVersion field not addressable")
  22. return reflect.NewAt(field.Type(), unsafe.Pointer(field.UnsafeAddr())).Elem().String()
  23. }
  24. func TestGetSharedReqClient_ForceHTTP2SeparatesCache(t *testing.T) {
  25. sharedReqClients = sync.Map{}
  26. base := reqClientOptions{
  27. ProxyURL: "http://proxy.local:8080",
  28. Timeout: time.Second,
  29. }
  30. clientDefault, err := getSharedReqClient(base)
  31. require.NoError(t, err)
  32. force := base
  33. force.ForceHTTP2 = true
  34. clientForce, err := getSharedReqClient(force)
  35. require.NoError(t, err)
  36. require.NotSame(t, clientDefault, clientForce)
  37. require.NotEqual(t, buildReqClientKey(base), buildReqClientKey(force))
  38. }
  39. func TestGetSharedReqClient_ReuseCachedClient(t *testing.T) {
  40. sharedReqClients = sync.Map{}
  41. opts := reqClientOptions{
  42. ProxyURL: "http://proxy.local:8080",
  43. Timeout: 2 * time.Second,
  44. }
  45. first, err := getSharedReqClient(opts)
  46. require.NoError(t, err)
  47. second, err := getSharedReqClient(opts)
  48. require.NoError(t, err)
  49. require.Same(t, first, second)
  50. }
  51. func TestGetSharedReqClient_IgnoresNonClientCache(t *testing.T) {
  52. sharedReqClients = sync.Map{}
  53. opts := reqClientOptions{
  54. ProxyURL: " http://proxy.local:8080 ",
  55. Timeout: 3 * time.Second,
  56. }
  57. key := buildReqClientKey(opts)
  58. sharedReqClients.Store(key, "invalid")
  59. client, err := getSharedReqClient(opts)
  60. require.NoError(t, err)
  61. require.NotNil(t, client)
  62. loaded, ok := sharedReqClients.Load(key)
  63. require.True(t, ok)
  64. require.IsType(t, "invalid", loaded)
  65. }
  66. func TestGetSharedReqClient_ImpersonateAndProxy(t *testing.T) {
  67. sharedReqClients = sync.Map{}
  68. opts := reqClientOptions{
  69. ProxyURL: " http://proxy.local:8080 ",
  70. Timeout: 4 * time.Second,
  71. Impersonate: true,
  72. }
  73. client, err := getSharedReqClient(opts)
  74. require.NoError(t, err)
  75. require.NotNil(t, client)
  76. require.Equal(t, "http://proxy.local:8080|4s|true|false", buildReqClientKey(opts))
  77. }
  78. func TestGetSharedReqClient_InvalidProxyURL(t *testing.T) {
  79. sharedReqClients = sync.Map{}
  80. opts := reqClientOptions{
  81. ProxyURL: "://missing-scheme",
  82. Timeout: time.Second,
  83. }
  84. _, err := getSharedReqClient(opts)
  85. require.Error(t, err)
  86. require.Contains(t, err.Error(), "invalid proxy URL")
  87. }
  88. func TestGetSharedReqClient_ProxyURLMissingHost(t *testing.T) {
  89. sharedReqClients = sync.Map{}
  90. opts := reqClientOptions{
  91. ProxyURL: "http://",
  92. Timeout: time.Second,
  93. }
  94. _, err := getSharedReqClient(opts)
  95. require.Error(t, err)
  96. require.Contains(t, err.Error(), "proxy URL missing host")
  97. }
  98. func TestCreateOpenAIReqClient_Timeout120Seconds(t *testing.T) {
  99. sharedReqClients = sync.Map{}
  100. client, err := createOpenAIReqClient("http://proxy.local:8080")
  101. require.NoError(t, err)
  102. require.Equal(t, 120*time.Second, client.GetClient().Timeout)
  103. }
  104. func TestCreateGeminiReqClient_ForceHTTP2Disabled(t *testing.T) {
  105. sharedReqClients = sync.Map{}
  106. client, err := createGeminiReqClient("http://proxy.local:8080")
  107. require.NoError(t, err)
  108. require.Equal(t, "", forceHTTPVersion(t, client))
  109. }
  110. func TestInstrumentReqClientRecordsDependency(t *testing.T) {
  111. server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
  112. w.WriteHeader(http.StatusNoContent)
  113. }))
  114. defer server.Close()
  115. collector := servertiming.New(time.Now())
  116. ctx := servertiming.WithCollector(context.Background(), collector)
  117. client := instrumentReqClient(req.C())
  118. response, err := client.R().SetContext(ctx).Get(server.URL)
  119. require.NoError(t, err)
  120. require.Equal(t, http.StatusNoContent, response.StatusCode)
  121. header := collector.HeaderValue(time.Now(), "bypass")
  122. require.True(t, strings.Contains(header, "dep_http;dur="), header)
  123. }