feat(websearch): settings UI overhaul and quota improvements
- Remove Priority field, auto load-balance by quota remaining - Replace QuotaRefreshInterval (daily/weekly/monthly) with SubscribedAt (subscription date, monthly lazy refresh via Redis TTL) - Add collapsible provider cards, API key show/copy, usage progress bar - Add test endpoint (POST /web-search-emulation/test) bypassing quota - Wire WebSearchManagerBuilder on startup (was never called before) - Fix nextMonthlyReset day-of-month overflow (Jan 31 → Feb 28) - Fix non-deterministic sort in selectByQuotaWeight - Map ProxyID in builder for provider-level proxy tracking - Fix frontend timezone drift in subscribed_at date picker - Fix provider deletion index shift for expandedProviders state
This commit is contained in:
erio
@@ -19,26 +19,18 @@ import (
|
||||
"github.com/redis/go-redis/v9"
|
||||
)
|
||||
|
||||
// Quota refresh interval constants.
|
||||
const (
|
||||
QuotaRefreshDaily = "daily"
|
||||
QuotaRefreshWeekly = "weekly"
|
||||
QuotaRefreshMonthly = "monthly"
|
||||
)
|
||||
|
||||
// ProviderConfig holds the configuration for a single search provider.
|
||||
type ProviderConfig struct {
|
||||
Type string `json:"type"` // ProviderTypeBrave | ProviderTypeTavily
|
||||
APIKey string `json:"api_key"` // secret
|
||||
Priority int `json:"priority"` // lower = higher priority
|
||||
QuotaLimit int64 `json:"quota_limit"` // 0 = unlimited
|
||||
QuotaRefreshInterval string `json:"quota_refresh_interval"` // QuotaRefreshDaily / Weekly / Monthly
|
||||
ProxyURL string `json:"-"` // resolved proxy URL (not persisted)
|
||||
ProxyID int64 `json:"-"` // resolved proxy ID for unavailability tracking
|
||||
ExpiresAt *int64 `json:"expires_at,omitempty"` // optional expiration (unix seconds)
|
||||
Type string `json:"type"` // ProviderTypeBrave | ProviderTypeTavily
|
||||
APIKey string `json:"api_key"` // secret
|
||||
QuotaLimit int64 `json:"quota_limit"` // 0 = unlimited
|
||||
SubscribedAt *int64 `json:"subscribed_at,omitempty"` // subscription start (unix seconds); quota resets monthly from this date
|
||||
ProxyURL string `json:"-"` // resolved proxy URL (not persisted)
|
||||
ProxyID int64 `json:"-"` // resolved proxy ID for unavailability tracking
|
||||
ExpiresAt *int64 `json:"expires_at,omitempty"` // optional expiration (unix seconds)
|
||||
}
|
||||
|
||||
// Manager selects providers by priority and tracks quota via Redis.
|
||||
// Manager selects providers by quota-weighted load balancing and tracks quota via Redis.
|
||||
type Manager struct {
|
||||
configs []ProviderConfig
|
||||
redis *redis.Client
|
||||
@@ -58,6 +50,7 @@ const (
|
||||
proxyUnavailableKey = "websearch:proxy_unavailable:%d"
|
||||
proxyUnavailableTTL = 5 * time.Minute
|
||||
quotaTTLBuffer = 24 * time.Hour
|
||||
defaultQuotaTTL = 31*24*time.Hour + quotaTTLBuffer // fallback when no subscription date
|
||||
maxCachedClients = 100
|
||||
)
|
||||
|
||||
@@ -80,14 +73,12 @@ return val
|
||||
`)
|
||||
|
||||
// NewManager creates a Manager with the given provider configs and Redis client.
|
||||
// Provider order is preserved as-is; selectByQuotaWeight handles load balancing.
|
||||
func NewManager(configs []ProviderConfig, redisClient *redis.Client) *Manager {
|
||||
sorted := make([]ProviderConfig, len(configs))
|
||||
copy(sorted, configs)
|
||||
sort.Slice(sorted, func(i, j int) bool {
|
||||
return sorted[i].Priority < sorted[j].Priority
|
||||
})
|
||||
copied := make([]ProviderConfig, len(configs))
|
||||
copy(copied, configs)
|
||||
return &Manager{
|
||||
configs: sorted,
|
||||
configs: copied,
|
||||
redis: redisClient,
|
||||
clientCache: make(map[string]*http.Client),
|
||||
}
|
||||
@@ -162,21 +153,28 @@ type weighted struct {
|
||||
// Providers with quota_limit=0 (no limit set) get weight 0 and are placed last.
|
||||
// Among providers with quota, higher remaining quota = higher priority.
|
||||
func (m *Manager) selectByQuotaWeight(ctx context.Context, candidates []ProviderConfig) []ProviderConfig {
|
||||
items := m.computeWeights(ctx, candidates)
|
||||
withQuota, withoutQuota := partitionByQuota(items)
|
||||
sortByStableRandomWeight(withQuota)
|
||||
return mergeWeightedResults(withQuota, withoutQuota, len(candidates))
|
||||
}
|
||||
|
||||
func (m *Manager) computeWeights(ctx context.Context, candidates []ProviderConfig) []weighted {
|
||||
items := make([]weighted, 0, len(candidates))
|
||||
for _, cfg := range candidates {
|
||||
w := int64(0)
|
||||
if cfg.QuotaLimit > 0 {
|
||||
used, _ := m.GetUsage(ctx, cfg.Type, cfg.QuotaRefreshInterval)
|
||||
remaining := cfg.QuotaLimit - used
|
||||
if remaining > 0 {
|
||||
used, _ := m.GetUsage(ctx, cfg.Type)
|
||||
if remaining := cfg.QuotaLimit - used; remaining > 0 {
|
||||
w = remaining
|
||||
}
|
||||
}
|
||||
items = append(items, weighted{cfg: cfg, weight: w})
|
||||
}
|
||||
return items
|
||||
}
|
||||
|
||||
// Separate providers with quota (weight > 0) from those without (weight == 0)
|
||||
var withQuota, withoutQuota []weighted
|
||||
func partitionByQuota(items []weighted) (withQuota, withoutQuota []weighted) {
|
||||
for _, item := range items {
|
||||
if item.weight > 0 {
|
||||
withQuota = append(withQuota, item)
|
||||
@@ -184,18 +182,26 @@ func (m *Manager) selectByQuotaWeight(ctx context.Context, candidates []Provider
|
||||
withoutQuota = append(withoutQuota, item)
|
||||
}
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
// Within quota group: weighted random sort (higher remaining = more likely first)
|
||||
if len(withQuota) > 1 {
|
||||
sort.Slice(withQuota, func(i, j int) bool {
|
||||
wi := float64(withQuota[i].weight) * (0.5 + rand.Float64())
|
||||
wj := float64(withQuota[j].weight) * (0.5 + rand.Float64())
|
||||
return wi > wj
|
||||
})
|
||||
// sortByStableRandomWeight assigns a fixed random factor to each item before sorting,
|
||||
// ensuring deterministic sort behavior (transitivity) within a single call.
|
||||
func sortByStableRandomWeight(items []weighted) {
|
||||
if len(items) <= 1 {
|
||||
return
|
||||
}
|
||||
factors := make([]float64, len(items))
|
||||
for i, item := range items {
|
||||
factors[i] = float64(item.weight) * (0.5 + rand.Float64())
|
||||
}
|
||||
sort.Slice(items, func(i, j int) bool {
|
||||
return factors[i] > factors[j]
|
||||
})
|
||||
}
|
||||
|
||||
// Build final order: quota providers first, then no-quota providers (original priority order)
|
||||
result := make([]ProviderConfig, 0, len(candidates))
|
||||
func mergeWeightedResults(withQuota, withoutQuota []weighted, capacity int) []ProviderConfig {
|
||||
result := make([]ProviderConfig, 0, capacity)
|
||||
for _, item := range withQuota {
|
||||
result = append(result, item.cfg)
|
||||
}
|
||||
@@ -294,8 +300,8 @@ func (m *Manager) tryReserveQuota(ctx context.Context, cfg ProviderConfig) (bool
|
||||
slog.Warn("websearch: Redis unavailable, quota check skipped", "provider", cfg.Type)
|
||||
return true, false
|
||||
}
|
||||
key := quotaRedisKey(cfg.Type, cfg.QuotaRefreshInterval)
|
||||
ttlSec := int(quotaTTL(cfg.QuotaRefreshInterval).Seconds())
|
||||
key := quotaRedisKey(cfg.Type)
|
||||
ttlSec := int(quotaTTLFromSubscription(cfg.SubscribedAt).Seconds())
|
||||
newVal, err := quotaIncrScript.Run(ctx, m.redis, []string{key}, ttlSec).Int64()
|
||||
if err != nil {
|
||||
slog.Warn("websearch: quota Lua INCR failed, allowing request",
|
||||
@@ -318,7 +324,7 @@ func (m *Manager) rollbackQuota(ctx context.Context, cfg ProviderConfig) {
|
||||
if cfg.QuotaLimit <= 0 || m.redis == nil {
|
||||
return
|
||||
}
|
||||
key := quotaRedisKey(cfg.Type, cfg.QuotaRefreshInterval)
|
||||
key := quotaRedisKey(cfg.Type)
|
||||
if err := m.redis.Decr(ctx, key).Err(); err != nil {
|
||||
slog.Warn("websearch: quota rollback DECR failed",
|
||||
"provider", cfg.Type, "error", err)
|
||||
@@ -327,6 +333,25 @@ func (m *Manager) rollbackQuota(ctx context.Context, cfg ProviderConfig) {
|
||||
|
||||
// --- Search execution ---
|
||||
|
||||
// TestSearch executes a search using the first available provider without reserving quota.
|
||||
// Intended for admin test functionality only.
|
||||
func (m *Manager) TestSearch(ctx context.Context, req SearchRequest) (*SearchResponse, string, error) {
|
||||
if strings.TrimSpace(req.Query) == "" {
|
||||
return nil, "", fmt.Errorf("websearch: empty search query")
|
||||
}
|
||||
for _, cfg := range m.configs {
|
||||
if !m.isProviderAvailable(cfg) {
|
||||
continue
|
||||
}
|
||||
resp, err := m.executeSearch(ctx, cfg, req)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
return resp, cfg.Type, nil
|
||||
}
|
||||
return nil, "", fmt.Errorf("websearch: no available provider")
|
||||
}
|
||||
|
||||
func (m *Manager) executeSearch(ctx context.Context, cfg ProviderConfig, req SearchRequest) (*SearchResponse, error) {
|
||||
proxyURL := cfg.ProxyURL
|
||||
if req.ProxyURL != "" {
|
||||
@@ -384,11 +409,11 @@ func newHTTPClient(proxyURL string) (*http.Client, error) {
|
||||
}
|
||||
|
||||
// GetUsage returns the current usage count for the given provider.
|
||||
func (m *Manager) GetUsage(ctx context.Context, providerType, refreshInterval string) (int64, error) {
|
||||
func (m *Manager) GetUsage(ctx context.Context, providerType string) (int64, error) {
|
||||
if m.redis == nil {
|
||||
return 0, nil
|
||||
}
|
||||
key := quotaRedisKey(providerType, refreshInterval)
|
||||
key := quotaRedisKey(providerType)
|
||||
val, err := m.redis.Get(ctx, key).Int64()
|
||||
if err == redis.Nil {
|
||||
return 0, nil
|
||||
@@ -400,7 +425,7 @@ func (m *Manager) GetUsage(ctx context.Context, providerType, refreshInterval st
|
||||
func (m *Manager) GetAllUsage(ctx context.Context) map[string]int64 {
|
||||
result := make(map[string]int64, len(m.configs))
|
||||
for _, cfg := range m.configs {
|
||||
used, _ := m.GetUsage(ctx, cfg.Type, cfg.QuotaRefreshInterval)
|
||||
used, _ := m.GetUsage(ctx, cfg.Type)
|
||||
result[cfg.Type] = used
|
||||
}
|
||||
return result
|
||||
@@ -423,30 +448,56 @@ func (m *Manager) buildProvider(cfg ProviderConfig, client *http.Client) Provide
|
||||
|
||||
// --- Redis key helpers ---
|
||||
|
||||
func quotaRedisKey(providerType, refreshInterval string) string {
|
||||
return quotaKeyPrefix + providerType + ":" + periodKey(refreshInterval)
|
||||
func quotaRedisKey(providerType string) string {
|
||||
return quotaKeyPrefix + providerType
|
||||
}
|
||||
|
||||
func periodKey(refreshInterval string) string {
|
||||
// quotaTTLFromSubscription calculates the TTL for the quota counter based on
|
||||
// the provider's subscription start date. Quota resets monthly from that date.
|
||||
// When the Redis key expires naturally, the next INCR creates a fresh counter (lazy refresh).
|
||||
func quotaTTLFromSubscription(subscribedAt *int64) time.Duration {
|
||||
if subscribedAt == nil || *subscribedAt == 0 {
|
||||
return defaultQuotaTTL
|
||||
}
|
||||
next := nextMonthlyReset(time.Unix(*subscribedAt, 0).UTC())
|
||||
ttl := time.Until(next) + quotaTTLBuffer
|
||||
if ttl <= quotaTTLBuffer {
|
||||
// Already past the reset — next cycle
|
||||
ttl = defaultQuotaTTL
|
||||
}
|
||||
return ttl
|
||||
}
|
||||
|
||||
// nextMonthlyReset returns the next monthly reset time based on the subscription start date.
|
||||
// E.g., subscribed on Jan 15 → resets on Feb 15, Mar 15, etc.
|
||||
// Handles day-of-month overflow: Jan 31 → Feb 28 (not Mar 3).
|
||||
func nextMonthlyReset(subscribedAt time.Time) time.Time {
|
||||
now := time.Now().UTC()
|
||||
switch refreshInterval {
|
||||
case QuotaRefreshDaily:
|
||||
return now.Format("2006-01-02")
|
||||
case QuotaRefreshWeekly:
|
||||
year, week := now.ISOWeek()
|
||||
return fmt.Sprintf("%d-W%02d", year, week)
|
||||
default:
|
||||
return now.Format("2006-01")
|
||||
if subscribedAt.IsZero() {
|
||||
return now.AddDate(0, 1, 0)
|
||||
}
|
||||
months := (now.Year()-subscribedAt.Year())*12 + int(now.Month()-subscribedAt.Month())
|
||||
if months < 0 {
|
||||
months = 0
|
||||
}
|
||||
candidate := addMonthsClamped(subscribedAt, months)
|
||||
if candidate.After(now) {
|
||||
return candidate
|
||||
}
|
||||
return addMonthsClamped(subscribedAt, months+1)
|
||||
}
|
||||
|
||||
func quotaTTL(refreshInterval string) time.Duration {
|
||||
switch refreshInterval {
|
||||
case QuotaRefreshDaily:
|
||||
return 24*time.Hour + quotaTTLBuffer
|
||||
case QuotaRefreshWeekly:
|
||||
return 7*24*time.Hour + quotaTTLBuffer
|
||||
default:
|
||||
return 31*24*time.Hour + quotaTTLBuffer
|
||||
// addMonthsClamped adds N months to a date, clamping the day to the last day of the target month.
|
||||
// E.g., Jan 31 + 1 month = Feb 28 (not Mar 3).
|
||||
func addMonthsClamped(t time.Time, months int) time.Time {
|
||||
y, m, d := t.Date()
|
||||
targetMonth := time.Month(int(m) + months)
|
||||
targetYear := y + int(targetMonth-1)/12
|
||||
targetMonth = (targetMonth-1)%12 + 1
|
||||
// Last day of the target month
|
||||
lastDay := time.Date(targetYear, targetMonth+1, 0, 0, 0, 0, 0, time.UTC).Day()
|
||||
if d > lastDay {
|
||||
d = lastDay
|
||||
}
|
||||
return time.Date(targetYear, targetMonth, d, 0, 0, 0, 0, time.UTC)
|
||||
}
|
||||
|
||||
@@ -12,14 +12,14 @@ import (
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestNewManager_SortsByPriority(t *testing.T) {
|
||||
func TestNewManager_PreservesOrder(t *testing.T) {
|
||||
configs := []ProviderConfig{
|
||||
{Type: "brave", APIKey: "k3", Priority: 30},
|
||||
{Type: "tavily", APIKey: "k1", Priority: 10},
|
||||
{Type: "brave", APIKey: "k3"},
|
||||
{Type: "tavily", APIKey: "k1"},
|
||||
}
|
||||
m := NewManager(configs, nil)
|
||||
require.Equal(t, 10, m.configs[0].Priority)
|
||||
require.Equal(t, 30, m.configs[1].Priority)
|
||||
require.Equal(t, "brave", m.configs[0].Type)
|
||||
require.Equal(t, "tavily", m.configs[1].Type)
|
||||
}
|
||||
|
||||
func TestManager_SearchWithBestProvider_EmptyQuery(t *testing.T) {
|
||||
@@ -46,8 +46,7 @@ func TestManager_SearchWithBestProvider_SkipExpired(t *testing.T) {
|
||||
require.ErrorContains(t, err, "no available provider")
|
||||
}
|
||||
|
||||
func TestManager_SearchWithBestProvider_PriorityOrder(t *testing.T) {
|
||||
// Create two mock servers that return different results
|
||||
func TestManager_SearchWithBestProvider_UsesFirstAvailable(t *testing.T) {
|
||||
srvBrave := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
resp := braveResponse{}
|
||||
resp.Web.Results = []braveResult{{URL: "https://brave.com", Title: "Brave", Description: "from brave"}}
|
||||
@@ -55,17 +54,15 @@ func TestManager_SearchWithBestProvider_PriorityOrder(t *testing.T) {
|
||||
}))
|
||||
defer srvBrave.Close()
|
||||
|
||||
// Override brave endpoint for test
|
||||
origURL := *braveSearchURL
|
||||
u, _ := http.NewRequest("GET", srvBrave.URL, nil)
|
||||
*braveSearchURL = *u.URL
|
||||
defer func() { *braveSearchURL = origURL }()
|
||||
|
||||
m := NewManager([]ProviderConfig{
|
||||
{Type: "brave", APIKey: "k1", Priority: 1},
|
||||
{Type: "tavily", APIKey: "k2", Priority: 2},
|
||||
{Type: "brave", APIKey: "k1"},
|
||||
{Type: "tavily", APIKey: "k2"},
|
||||
}, nil)
|
||||
// Inject the test server's client
|
||||
m.clientCache[srvBrave.URL] = srvBrave.Client()
|
||||
m.clientCache[""] = srvBrave.Client()
|
||||
|
||||
@@ -77,7 +74,6 @@ func TestManager_SearchWithBestProvider_PriorityOrder(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestManager_SearchWithBestProvider_NilRedis(t *testing.T) {
|
||||
// With nil Redis, quota check is skipped (always allowed)
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
|
||||
resp := braveResponse{}
|
||||
resp.Web.Results = []braveResult{{URL: "https://test.com", Title: "Test", Description: "result"}}
|
||||
@@ -91,8 +87,8 @@ func TestManager_SearchWithBestProvider_NilRedis(t *testing.T) {
|
||||
defer func() { *braveSearchURL = origURL }()
|
||||
|
||||
m := NewManager([]ProviderConfig{
|
||||
{Type: "brave", APIKey: "k", Priority: 1, QuotaLimit: 100},
|
||||
}, nil) // nil Redis
|
||||
{Type: "brave", APIKey: "k", QuotaLimit: 100},
|
||||
}, nil)
|
||||
m.clientCache[""] = srv.Client()
|
||||
|
||||
resp, _, err := m.SearchWithBestProvider(context.Background(), SearchRequest{Query: "test"})
|
||||
@@ -102,51 +98,98 @@ func TestManager_SearchWithBestProvider_NilRedis(t *testing.T) {
|
||||
|
||||
func TestManager_GetUsage_NilRedis(t *testing.T) {
|
||||
m := NewManager(nil, nil)
|
||||
used, err := m.GetUsage(context.Background(), "brave", "monthly")
|
||||
used, err := m.GetUsage(context.Background(), "brave")
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(0), used)
|
||||
}
|
||||
|
||||
func TestManager_GetAllUsage_NilRedis(t *testing.T) {
|
||||
m := NewManager([]ProviderConfig{
|
||||
{Type: "brave", QuotaRefreshInterval: "monthly"},
|
||||
{Type: "brave"},
|
||||
}, nil)
|
||||
usage := m.GetAllUsage(context.Background())
|
||||
require.Equal(t, int64(0), usage["brave"])
|
||||
}
|
||||
|
||||
// --- Key/TTL helpers ---
|
||||
// --- Quota TTL from subscription ---
|
||||
|
||||
func TestQuotaTTL_Daily(t *testing.T) {
|
||||
require.Equal(t, 24*time.Hour+quotaTTLBuffer, quotaTTL(QuotaRefreshDaily))
|
||||
func TestQuotaTTLFromSubscription_NilSubscription(t *testing.T) {
|
||||
ttl := quotaTTLFromSubscription(nil)
|
||||
require.Equal(t, defaultQuotaTTL, ttl)
|
||||
}
|
||||
|
||||
func TestQuotaTTL_Weekly(t *testing.T) {
|
||||
require.Equal(t, 7*24*time.Hour+quotaTTLBuffer, quotaTTL(QuotaRefreshWeekly))
|
||||
func TestQuotaTTLFromSubscription_ZeroSubscription(t *testing.T) {
|
||||
zero := int64(0)
|
||||
ttl := quotaTTLFromSubscription(&zero)
|
||||
require.Equal(t, defaultQuotaTTL, ttl)
|
||||
}
|
||||
|
||||
func TestQuotaTTL_Monthly(t *testing.T) {
|
||||
require.Equal(t, 31*24*time.Hour+quotaTTLBuffer, quotaTTL(QuotaRefreshMonthly))
|
||||
func TestQuotaTTLFromSubscription_ValidSubscription(t *testing.T) {
|
||||
// Subscribed 10 days ago — next reset in ~20 days
|
||||
sub := time.Now().Add(-10 * 24 * time.Hour).Unix()
|
||||
ttl := quotaTTLFromSubscription(&sub)
|
||||
require.Greater(t, ttl, 15*24*time.Hour) // at least 15 days
|
||||
require.Less(t, ttl, 25*24*time.Hour+quotaTTLBuffer)
|
||||
}
|
||||
|
||||
func TestPeriodKey_Daily(t *testing.T) {
|
||||
key := periodKey(QuotaRefreshDaily)
|
||||
require.Regexp(t, `^\d{4}-\d{2}-\d{2}$`, key)
|
||||
func TestNextMonthlyReset_SubscribedRecentPast(t *testing.T) {
|
||||
// Subscribed on the 10th of this month (always valid day)
|
||||
now := time.Now().UTC()
|
||||
sub := time.Date(now.Year(), now.Month(), 10, 0, 0, 0, 0, time.UTC)
|
||||
next := nextMonthlyReset(sub)
|
||||
require.True(t, next.After(now) || next.Equal(now), "next reset should be in the future or now")
|
||||
require.True(t, next.Before(now.AddDate(0, 1, 1)))
|
||||
}
|
||||
|
||||
func TestPeriodKey_Weekly(t *testing.T) {
|
||||
key := periodKey(QuotaRefreshWeekly)
|
||||
require.Regexp(t, `^\d{4}-W\d{2}$`, key)
|
||||
func TestNextMonthlyReset_SubscribedLongAgo(t *testing.T) {
|
||||
// Subscribed 6 months ago on the 1st
|
||||
sub := time.Now().UTC().AddDate(0, -6, 0)
|
||||
sub = time.Date(sub.Year(), sub.Month(), 1, 0, 0, 0, 0, time.UTC)
|
||||
next := nextMonthlyReset(sub)
|
||||
require.True(t, next.After(time.Now().UTC()))
|
||||
// Should be within the next 31 days
|
||||
require.True(t, next.Before(time.Now().UTC().AddDate(0, 1, 1)))
|
||||
}
|
||||
|
||||
func TestPeriodKey_Monthly(t *testing.T) {
|
||||
key := periodKey(QuotaRefreshMonthly)
|
||||
require.Regexp(t, `^\d{4}-\d{2}$`, key)
|
||||
func TestNextMonthlyReset_FutureSubscription(t *testing.T) {
|
||||
sub := time.Now().UTC().AddDate(0, 0, 5)
|
||||
next := nextMonthlyReset(sub)
|
||||
require.True(t, next.After(time.Now().UTC()))
|
||||
}
|
||||
|
||||
func TestAddMonthsClamped_Jan31ToFeb(t *testing.T) {
|
||||
sub := time.Date(2026, 1, 31, 0, 0, 0, 0, time.UTC)
|
||||
next := addMonthsClamped(sub, 1)
|
||||
require.Equal(t, time.Month(2), next.Month())
|
||||
require.Equal(t, 28, next.Day()) // Feb 28 (2026 is not a leap year)
|
||||
}
|
||||
|
||||
func TestAddMonthsClamped_Jan31ToFebLeapYear(t *testing.T) {
|
||||
sub := time.Date(2028, 1, 31, 0, 0, 0, 0, time.UTC)
|
||||
next := addMonthsClamped(sub, 1)
|
||||
require.Equal(t, time.Month(2), next.Month())
|
||||
require.Equal(t, 29, next.Day()) // Feb 29 (2028 is a leap year)
|
||||
}
|
||||
|
||||
func TestAddMonthsClamped_Mar31ToApr(t *testing.T) {
|
||||
sub := time.Date(2026, 3, 31, 0, 0, 0, 0, time.UTC)
|
||||
next := addMonthsClamped(sub, 1)
|
||||
require.Equal(t, time.Month(4), next.Month())
|
||||
require.Equal(t, 30, next.Day()) // Apr has 30 days
|
||||
}
|
||||
|
||||
func TestAddMonthsClamped_NormalDay(t *testing.T) {
|
||||
sub := time.Date(2026, 1, 15, 0, 0, 0, 0, time.UTC)
|
||||
next := addMonthsClamped(sub, 1)
|
||||
require.Equal(t, time.Month(2), next.Month())
|
||||
require.Equal(t, 15, next.Day()) // no clamping needed
|
||||
}
|
||||
|
||||
// --- Redis key ---
|
||||
|
||||
func TestQuotaRedisKey_Format(t *testing.T) {
|
||||
key := quotaRedisKey("brave", QuotaRefreshDaily)
|
||||
require.Contains(t, key, "websearch:quota:brave:")
|
||||
key := quotaRedisKey("brave")
|
||||
require.Equal(t, "websearch:quota:brave", key)
|
||||
}
|
||||
|
||||
// --- isProviderAvailable ---
|
||||
@@ -173,9 +216,7 @@ func TestIsProviderAvailable_Valid(t *testing.T) {
|
||||
|
||||
func TestResolveProxyID_AccountProxyOverrides(t *testing.T) {
|
||||
cfg := ProviderConfig{ProxyID: 42}
|
||||
// account proxy present → return 0 (account proxy has no config-level ID)
|
||||
require.Equal(t, int64(0), resolveProxyID(cfg, "http://account-proxy:8080"))
|
||||
// no account proxy → return provider's proxy ID
|
||||
require.Equal(t, int64(42), resolveProxyID(cfg, ""))
|
||||
}
|
||||
|
||||
@@ -186,28 +227,23 @@ func TestIsProxyError_Nil(t *testing.T) {
|
||||
}
|
||||
|
||||
func TestIsProxyError_ConnectionRefused(t *testing.T) {
|
||||
err := fmt.Errorf("dial tcp: connection refused")
|
||||
require.True(t, isProxyError(err))
|
||||
require.True(t, isProxyError(fmt.Errorf("dial tcp: connection refused")))
|
||||
}
|
||||
|
||||
func TestIsProxyError_Timeout(t *testing.T) {
|
||||
err := fmt.Errorf("i/o timeout while connecting to proxy")
|
||||
require.True(t, isProxyError(err))
|
||||
require.True(t, isProxyError(fmt.Errorf("i/o timeout while connecting to proxy")))
|
||||
}
|
||||
|
||||
func TestIsProxyError_SOCKS(t *testing.T) {
|
||||
err := fmt.Errorf("socks connect failed")
|
||||
require.True(t, isProxyError(err))
|
||||
require.True(t, isProxyError(fmt.Errorf("socks connect failed")))
|
||||
}
|
||||
|
||||
func TestIsProxyError_TLSHandshake(t *testing.T) {
|
||||
err := fmt.Errorf("tls handshake timeout")
|
||||
require.True(t, isProxyError(err))
|
||||
require.True(t, isProxyError(fmt.Errorf("tls handshake timeout")))
|
||||
}
|
||||
|
||||
func TestIsProxyError_APIError_NotProxy(t *testing.T) {
|
||||
err := fmt.Errorf("API rate limit exceeded")
|
||||
require.False(t, isProxyError(err))
|
||||
require.False(t, isProxyError(fmt.Errorf("API rate limit exceeded")))
|
||||
}
|
||||
|
||||
// --- isProxyAvailable (nil Redis) ---
|
||||
@@ -225,14 +261,13 @@ func TestIsProxyAvailable_ZeroID(t *testing.T) {
|
||||
// --- selectByQuotaWeight ---
|
||||
|
||||
func TestSelectByQuotaWeight_NoQuotaLast(t *testing.T) {
|
||||
m := NewManager(nil, nil) // nil Redis → GetUsage returns 0
|
||||
m := NewManager(nil, nil)
|
||||
candidates := []ProviderConfig{
|
||||
{Type: "brave", APIKey: "k1", QuotaLimit: 0}, // no limit → weight 0
|
||||
{Type: "tavily", APIKey: "k2", QuotaLimit: 100}, // remaining 100
|
||||
{Type: "brave", APIKey: "k1", QuotaLimit: 0},
|
||||
{Type: "tavily", APIKey: "k2", QuotaLimit: 100},
|
||||
}
|
||||
result := m.selectByQuotaWeight(context.Background(), candidates)
|
||||
require.Len(t, result, 2)
|
||||
// tavily (with quota) should come first
|
||||
require.Equal(t, "tavily", result[0].Type)
|
||||
require.Equal(t, "brave", result[1].Type)
|
||||
}
|
||||
@@ -245,7 +280,6 @@ func TestSelectByQuotaWeight_AllNoQuota(t *testing.T) {
|
||||
}
|
||||
result := m.selectByQuotaWeight(context.Background(), candidates)
|
||||
require.Len(t, result, 2)
|
||||
// both have weight 0, original order preserved
|
||||
}
|
||||
|
||||
func TestSelectByQuotaWeight_Empty(t *testing.T) {
|
||||
|
||||
Reference in New Issue
Block a user