104 lines
4.2 KiB
Go
104 lines
4.2 KiB
Go
package health
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"proxy-pool/internal/config"
|
|
"proxy-pool/internal/domain/activitypool"
|
|
)
|
|
|
|
type unhealthyRemoverStub struct {
|
|
command activitypool.UnhealthySweepCommand
|
|
calls int
|
|
result activitypool.UnhealthySweepResult
|
|
drainCalls int
|
|
drainNow time.Time
|
|
candidates []activitypool.UnhealthyDrainCandidate
|
|
}
|
|
|
|
func (store *unhealthyRemoverStub) BeginUnhealthyDrain(
|
|
_ context.Context,
|
|
now time.Time,
|
|
candidate activitypool.UnhealthyDrainCandidate,
|
|
) (bool, error) {
|
|
store.drainCalls++
|
|
store.drainNow = now
|
|
store.candidates = append(store.candidates, candidate)
|
|
return true, nil
|
|
}
|
|
|
|
func (store *unhealthyRemoverStub) SweepUnhealthy(
|
|
_ context.Context,
|
|
command activitypool.UnhealthySweepCommand,
|
|
) (activitypool.UnhealthySweepResult, error) {
|
|
store.calls++
|
|
store.command = command
|
|
return store.result, nil
|
|
}
|
|
|
|
func TestConfiguredUnhealthyReaperUsesEffectiveEnabledPolicies(t *testing.T) {
|
|
now := time.Date(2026, 8, 2, 14, 0, 0, 0, time.UTC)
|
|
store := &unhealthyRemoverStub{result: activitypool.UnhealthySweepResult{Removed: 2}}
|
|
reaper, err := NewConfiguredUnhealthyReaper(&configurationSourceStub{configuration: &config.Config{
|
|
Defaults: config.Defaults{Check: config.Check{UnhealthyRemoveAfter: config.Duration(2 * time.Minute)}},
|
|
Upstreams: map[string]config.Upstream{
|
|
"provider-a": {Enabled: true, Check: config.Check{UnhealthyRemoveAfter: config.Duration(3 * time.Minute)}},
|
|
"provider-b": {Enabled: true},
|
|
"provider-c": {Enabled: false},
|
|
},
|
|
}}, store, UnhealthyReaperOptions{PollInterval: time.Second, BatchSize: 32, Now: func() time.Time { return now }})
|
|
if err != nil {
|
|
t.Fatalf("NewConfiguredUnhealthyReaper() = %v", err)
|
|
}
|
|
result, err := reaper.Tick(context.Background())
|
|
if err != nil || result.Removed != 2 || result.DeferredOwned != 0 || len(result.DrainCandidates) != 0 {
|
|
t.Fatalf("Tick() = %+v, %v", result, err)
|
|
}
|
|
if store.calls != 1 || !store.command.Now.Equal(now) || store.command.Limit != 32 ||
|
|
store.command.RemoveAfterByUpstream["provider-a"] != 3*time.Minute ||
|
|
store.command.RemoveAfterByUpstream["provider-b"] != 2*time.Minute ||
|
|
len(store.command.RemoveAfterByUpstream) != 2 {
|
|
t.Fatalf("SweepUnhealthy() command = %+v", store.command)
|
|
}
|
|
}
|
|
|
|
func TestConfiguredUnhealthyReaperSkipsDisabledRemoval(t *testing.T) {
|
|
store := &unhealthyRemoverStub{}
|
|
reaper, err := NewConfiguredUnhealthyReaper(&configurationSourceStub{configuration: &config.Config{
|
|
Upstreams: map[string]config.Upstream{"provider-a": {Enabled: true}},
|
|
}}, store, UnhealthyReaperOptions{PollInterval: time.Second, BatchSize: 1, Now: time.Now})
|
|
if err != nil {
|
|
t.Fatalf("NewConfiguredUnhealthyReaper() = %v", err)
|
|
}
|
|
if result, err := reaper.Tick(context.Background()); err != nil || result.Removed != 0 || result.DeferredOwned != 0 || len(result.DrainCandidates) != 0 || store.calls != 0 {
|
|
t.Fatalf("Tick(disabled) = (%+v, %v), calls=%d", result, err, store.calls)
|
|
}
|
|
}
|
|
|
|
func TestConfiguredUnhealthyReaperStartsBoundedDrainCandidates(t *testing.T) {
|
|
now := time.Date(2026, 8, 2, 14, 0, 0, 0, time.UTC)
|
|
candidate := activitypool.UnhealthyDrainCandidate{
|
|
ProxyID: "proxy-a", WorkerID: "worker-a", AssignmentEpoch: 7, UnhealthySince: now.Add(-time.Minute),
|
|
}
|
|
store := &unhealthyRemoverStub{result: activitypool.UnhealthySweepResult{
|
|
DeferredOwned: 1, DrainCandidates: []activitypool.UnhealthyDrainCandidate{candidate},
|
|
}}
|
|
reaper, err := NewConfiguredUnhealthyReaper(&configurationSourceStub{configuration: &config.Config{
|
|
Upstreams: map[string]config.Upstream{
|
|
"provider-a": {Enabled: true, Check: config.Check{UnhealthyRemoveAfter: config.Duration(time.Minute)}},
|
|
},
|
|
}}, store, UnhealthyReaperOptions{PollInterval: time.Second, BatchSize: 4, Now: func() time.Time { return now }})
|
|
if err != nil {
|
|
t.Fatalf("NewConfiguredUnhealthyReaper() = %v", err)
|
|
}
|
|
result, err := reaper.Tick(context.Background())
|
|
if err != nil || result.DeferredOwned != 1 || len(result.DrainCandidates) != 1 {
|
|
t.Fatalf("Tick() = %+v, %v", result, err)
|
|
}
|
|
if store.drainCalls != 1 || !store.drainNow.Equal(now) || len(store.candidates) != 1 || store.candidates[0] != candidate {
|
|
t.Fatalf("BeginUnhealthyDrain() calls=%d now=%v candidates=%+v", store.drainCalls, store.drainNow, store.candidates)
|
|
}
|
|
}
|