manager_test.go 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323
  1. package websearch
  2. import (
  3. "context"
  4. "encoding/json"
  5. "fmt"
  6. "net/http"
  7. "net/http/httptest"
  8. "testing"
  9. "time"
  10. "github.com/stretchr/testify/require"
  11. )
  12. func TestNewManager_PreservesOrder(t *testing.T) {
  13. configs := []ProviderConfig{
  14. {Type: "brave", APIKey: "k3"},
  15. {Type: "tavily", APIKey: "k1"},
  16. }
  17. m := NewManager(configs, nil)
  18. require.Equal(t, "brave", m.configs[0].Type)
  19. require.Equal(t, "tavily", m.configs[1].Type)
  20. }
  21. func TestManager_SearchWithBestProvider_EmptyQuery(t *testing.T) {
  22. m := NewManager([]ProviderConfig{{Type: "brave", APIKey: "k"}}, nil)
  23. _, _, err := m.SearchWithBestProvider(context.Background(), SearchRequest{Query: ""})
  24. require.ErrorContains(t, err, "empty search query")
  25. _, _, err = m.SearchWithBestProvider(context.Background(), SearchRequest{Query: " "})
  26. require.ErrorContains(t, err, "empty search query")
  27. }
  28. func TestManager_SearchWithBestProvider_SkipEmptyAPIKey(t *testing.T) {
  29. m := NewManager([]ProviderConfig{{Type: "brave", APIKey: ""}}, nil)
  30. _, _, err := m.SearchWithBestProvider(context.Background(), SearchRequest{Query: "test"})
  31. require.ErrorContains(t, err, "no available provider")
  32. }
  33. func TestManager_SearchWithBestProvider_SkipExpired(t *testing.T) {
  34. past := time.Now().Add(-1 * time.Hour).Unix()
  35. m := NewManager([]ProviderConfig{
  36. {Type: "brave", APIKey: "k", ExpiresAt: &past},
  37. }, nil)
  38. _, _, err := m.SearchWithBestProvider(context.Background(), SearchRequest{Query: "test"})
  39. require.ErrorContains(t, err, "no available provider")
  40. }
  41. func TestManager_SearchWithBestProvider_UsesFirstAvailable(t *testing.T) {
  42. srvBrave := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
  43. resp := braveResponse{}
  44. resp.Web.Results = []braveResult{{URL: "https://brave.com", Title: "Brave", Description: "from brave"}}
  45. _ = json.NewEncoder(w).Encode(resp)
  46. }))
  47. defer srvBrave.Close()
  48. origURL := *braveSearchURL
  49. u, _ := http.NewRequest("GET", srvBrave.URL, nil)
  50. *braveSearchURL = *u.URL
  51. defer func() { *braveSearchURL = origURL }()
  52. m := NewManager([]ProviderConfig{
  53. {Type: "brave", APIKey: "k1"},
  54. {Type: "tavily", APIKey: "k2"},
  55. }, nil)
  56. m.clientCache[srvBrave.URL] = srvBrave.Client()
  57. m.clientCache[""] = srvBrave.Client()
  58. resp, providerName, err := m.SearchWithBestProvider(context.Background(), SearchRequest{Query: "test"})
  59. require.NoError(t, err)
  60. require.Equal(t, "brave", providerName)
  61. require.Len(t, resp.Results, 1)
  62. require.Equal(t, "from brave", resp.Results[0].Snippet)
  63. }
  64. func TestManager_SearchWithBestProvider_NilRedis(t *testing.T) {
  65. srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
  66. resp := braveResponse{}
  67. resp.Web.Results = []braveResult{{URL: "https://test.com", Title: "Test", Description: "result"}}
  68. _ = json.NewEncoder(w).Encode(resp)
  69. }))
  70. defer srv.Close()
  71. origURL := *braveSearchURL
  72. u, _ := http.NewRequest("GET", srv.URL, nil)
  73. *braveSearchURL = *u.URL
  74. defer func() { *braveSearchURL = origURL }()
  75. m := NewManager([]ProviderConfig{
  76. {Type: "brave", APIKey: "k", QuotaLimit: 100},
  77. }, nil)
  78. m.clientCache[""] = srv.Client()
  79. resp, _, err := m.SearchWithBestProvider(context.Background(), SearchRequest{Query: "test"})
  80. require.NoError(t, err)
  81. require.Len(t, resp.Results, 1)
  82. }
  83. func TestManager_GetUsage_NilRedis(t *testing.T) {
  84. m := NewManager(nil, nil)
  85. used, err := m.GetUsage(context.Background(), "brave")
  86. require.NoError(t, err)
  87. require.Equal(t, int64(0), used)
  88. }
  89. func TestManager_GetAllUsage_NilRedis(t *testing.T) {
  90. m := NewManager([]ProviderConfig{
  91. {Type: "brave"},
  92. }, nil)
  93. usage := m.GetAllUsage(context.Background())
  94. require.Equal(t, int64(0), usage["brave"])
  95. }
  96. // --- Quota TTL from subscription ---
  97. func TestQuotaTTLFromSubscription_NilSubscription(t *testing.T) {
  98. ttl := quotaTTLFromSubscription(nil)
  99. require.Equal(t, defaultQuotaTTL, ttl)
  100. }
  101. func TestQuotaTTLFromSubscription_ZeroSubscription(t *testing.T) {
  102. zero := int64(0)
  103. ttl := quotaTTLFromSubscription(&zero)
  104. require.Equal(t, defaultQuotaTTL, ttl)
  105. }
  106. func TestQuotaTTLFromSubscription_ValidSubscription(t *testing.T) {
  107. // Subscribed 10 days ago — next reset in ~20 days
  108. sub := time.Now().Add(-10 * 24 * time.Hour).Unix()
  109. ttl := quotaTTLFromSubscription(&sub)
  110. require.Greater(t, ttl, 15*24*time.Hour) // at least 15 days
  111. require.Less(t, ttl, 25*24*time.Hour+quotaTTLBuffer)
  112. }
  113. func TestNextMonthlyReset_SubscribedRecentPast(t *testing.T) {
  114. // Subscribed on the 10th of this month (always valid day)
  115. now := time.Now().UTC()
  116. sub := time.Date(now.Year(), now.Month(), 10, 0, 0, 0, 0, time.UTC)
  117. next := nextMonthlyReset(sub)
  118. require.True(t, next.After(now) || next.Equal(now), "next reset should be in the future or now")
  119. require.True(t, next.Before(now.AddDate(0, 1, 1)))
  120. }
  121. func TestNextMonthlyReset_SubscribedLongAgo(t *testing.T) {
  122. // Subscribed 6 months ago on the 1st
  123. sub := time.Now().UTC().AddDate(0, -6, 0)
  124. sub = time.Date(sub.Year(), sub.Month(), 1, 0, 0, 0, 0, time.UTC)
  125. next := nextMonthlyReset(sub)
  126. require.True(t, next.After(time.Now().UTC()))
  127. // Should be within the next 31 days
  128. require.True(t, next.Before(time.Now().UTC().AddDate(0, 1, 1)))
  129. }
  130. func TestNextMonthlyReset_FutureSubscription(t *testing.T) {
  131. sub := time.Now().UTC().AddDate(0, 0, 5)
  132. next := nextMonthlyReset(sub)
  133. require.True(t, next.After(time.Now().UTC()))
  134. }
  135. func TestAddMonthsClamped_Jan31ToFeb(t *testing.T) {
  136. sub := time.Date(2026, 1, 31, 0, 0, 0, 0, time.UTC)
  137. next := addMonthsClamped(sub, 1)
  138. require.Equal(t, time.Month(2), next.Month())
  139. require.Equal(t, 28, next.Day()) // Feb 28 (2026 is not a leap year)
  140. }
  141. func TestAddMonthsClamped_Jan31ToFebLeapYear(t *testing.T) {
  142. sub := time.Date(2028, 1, 31, 0, 0, 0, 0, time.UTC)
  143. next := addMonthsClamped(sub, 1)
  144. require.Equal(t, time.Month(2), next.Month())
  145. require.Equal(t, 29, next.Day()) // Feb 29 (2028 is a leap year)
  146. }
  147. func TestAddMonthsClamped_Mar31ToApr(t *testing.T) {
  148. sub := time.Date(2026, 3, 31, 0, 0, 0, 0, time.UTC)
  149. next := addMonthsClamped(sub, 1)
  150. require.Equal(t, time.Month(4), next.Month())
  151. require.Equal(t, 30, next.Day()) // Apr has 30 days
  152. }
  153. func TestAddMonthsClamped_NormalDay(t *testing.T) {
  154. sub := time.Date(2026, 1, 15, 0, 0, 0, 0, time.UTC)
  155. next := addMonthsClamped(sub, 1)
  156. require.Equal(t, time.Month(2), next.Month())
  157. require.Equal(t, 15, next.Day()) // no clamping needed
  158. }
  159. // --- Redis key ---
  160. func TestQuotaRedisKey_Format(t *testing.T) {
  161. key := quotaRedisKey("brave")
  162. require.Equal(t, "websearch:quota:brave", key)
  163. }
  164. // --- isProviderAvailable ---
  165. func TestIsProviderAvailable_EmptyAPIKey(t *testing.T) {
  166. m := NewManager(nil, nil)
  167. require.False(t, m.isProviderAvailable(ProviderConfig{APIKey: ""}))
  168. }
  169. func TestIsProviderAvailable_Expired(t *testing.T) {
  170. m := NewManager(nil, nil)
  171. past := time.Now().Add(-1 * time.Hour).Unix()
  172. require.False(t, m.isProviderAvailable(ProviderConfig{APIKey: "k", ExpiresAt: &past}))
  173. }
  174. func TestIsProviderAvailable_Valid(t *testing.T) {
  175. m := NewManager(nil, nil)
  176. future := time.Now().Add(1 * time.Hour).Unix()
  177. require.True(t, m.isProviderAvailable(ProviderConfig{APIKey: "k", ExpiresAt: &future}))
  178. require.True(t, m.isProviderAvailable(ProviderConfig{APIKey: "k"})) // no expiry
  179. }
  180. // --- resolveProxyID ---
  181. func TestResolveProxyID_AccountProxyOverrides(t *testing.T) {
  182. cfg := ProviderConfig{ProxyID: 42}
  183. require.Equal(t, int64(0), resolveProxyID(cfg, "http://account-proxy:8080"))
  184. require.Equal(t, int64(42), resolveProxyID(cfg, ""))
  185. }
  186. // --- isProxyError ---
  187. func TestIsProxyError_Nil(t *testing.T) {
  188. require.False(t, isProxyError(nil))
  189. }
  190. func TestIsProxyError_ConnectionRefused(t *testing.T) {
  191. require.True(t, isProxyError(fmt.Errorf("dial tcp: connection refused")))
  192. }
  193. func TestIsProxyError_Timeout(t *testing.T) {
  194. require.True(t, isProxyError(fmt.Errorf("i/o timeout while connecting to proxy")))
  195. }
  196. func TestIsProxyError_SOCKS(t *testing.T) {
  197. require.True(t, isProxyError(fmt.Errorf("socks connect failed")))
  198. }
  199. func TestIsProxyError_TLSHandshake(t *testing.T) {
  200. require.True(t, isProxyError(fmt.Errorf("tls handshake timeout")))
  201. }
  202. func TestIsProxyError_APIError_NotProxy(t *testing.T) {
  203. require.False(t, isProxyError(fmt.Errorf("API rate limit exceeded")))
  204. }
  205. // --- isProxyAvailable (nil Redis) ---
  206. func TestIsProxyAvailable_NilRedis(t *testing.T) {
  207. m := NewManager(nil, nil)
  208. require.True(t, m.isProxyAvailable(context.Background(), 42))
  209. }
  210. func TestIsProxyAvailable_ZeroID(t *testing.T) {
  211. m := NewManager(nil, nil)
  212. require.True(t, m.isProxyAvailable(context.Background(), 0))
  213. }
  214. // --- selectByQuotaWeight ---
  215. func TestSelectByQuotaWeight_NoQuotaLast(t *testing.T) {
  216. m := NewManager(nil, nil)
  217. candidates := []ProviderConfig{
  218. {Type: "brave", APIKey: "k1", QuotaLimit: 0},
  219. {Type: "tavily", APIKey: "k2", QuotaLimit: 100},
  220. }
  221. result := m.selectByQuotaWeight(context.Background(), candidates)
  222. require.Len(t, result, 2)
  223. require.Equal(t, "tavily", result[0].Type)
  224. require.Equal(t, "brave", result[1].Type)
  225. }
  226. func TestSelectByQuotaWeight_AllNoQuota(t *testing.T) {
  227. m := NewManager(nil, nil)
  228. candidates := []ProviderConfig{
  229. {Type: "brave", APIKey: "k1", QuotaLimit: 0},
  230. {Type: "tavily", APIKey: "k2", QuotaLimit: 0},
  231. }
  232. result := m.selectByQuotaWeight(context.Background(), candidates)
  233. require.Len(t, result, 2)
  234. }
  235. func TestSelectByQuotaWeight_Empty(t *testing.T) {
  236. m := NewManager(nil, nil)
  237. result := m.selectByQuotaWeight(context.Background(), nil)
  238. require.Empty(t, result)
  239. }
  240. // --- newHTTPClient ---
  241. func TestNewHTTPClient_NoProxy(t *testing.T) {
  242. c, err := newHTTPClient("")
  243. require.NoError(t, err)
  244. require.NotNil(t, c)
  245. }
  246. func TestNewHTTPClient_InvalidProxy(t *testing.T) {
  247. _, err := newHTTPClient("://bad-url")
  248. require.Error(t, err)
  249. require.Contains(t, err.Error(), "invalid proxy URL")
  250. }
  251. func TestNewHTTPClient_ValidHTTPProxy(t *testing.T) {
  252. c, err := newHTTPClient("http://proxy.example.com:8080")
  253. require.NoError(t, err)
  254. require.NotNil(t, c)
  255. }
  256. func TestNewHTTPClient_ValidSOCKS5Proxy(t *testing.T) {
  257. c, err := newHTTPClient("socks5://proxy.example.com:1080")
  258. require.NoError(t, err)
  259. require.NotNil(t, c)
  260. }
  261. // --- ResetUsage ---
  262. func TestManager_ResetUsage_NilRedis(t *testing.T) {
  263. m := NewManager(nil, nil)
  264. err := m.ResetUsage(context.Background(), "brave")
  265. require.NoError(t, err)
  266. }