115 lines
4.4 KiB
Go
115 lines
4.4 KiB
Go
package health
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"proxy-pool/internal/config"
|
|
"proxy-pool/internal/domain/activitypool"
|
|
"proxy-pool/internal/domain/adminstate"
|
|
proxyDomain "proxy-pool/internal/domain/proxy"
|
|
)
|
|
|
|
type upstreamDrainPolicySourceStub struct {
|
|
policies []activitypool.UpstreamDrainPolicy
|
|
err error
|
|
}
|
|
|
|
func (source upstreamDrainPolicySourceStub) ReadUpstreamDrainPolicies(context.Context) ([]activitypool.UpstreamDrainPolicy, error) {
|
|
return append([]activitypool.UpstreamDrainPolicy(nil), source.policies...), source.err
|
|
}
|
|
|
|
type upstreamDrainConfigurationStub struct {
|
|
configuration *config.Config
|
|
revision uint64
|
|
}
|
|
|
|
type upstreamDrainStateStub struct{ snapshot adminstate.Snapshot }
|
|
|
|
func (source upstreamDrainStateStub) Snapshot(context.Context) (adminstate.Snapshot, error) {
|
|
return source.snapshot, nil
|
|
}
|
|
|
|
func (source upstreamDrainConfigurationStub) Snapshot() (*config.Config, uint64) {
|
|
return source.configuration, source.revision
|
|
}
|
|
|
|
func TestConfiguredUpstreamDrainReaperStartsBoundedDisabledAssignments(t *testing.T) {
|
|
now := time.Date(2026, 8, 2, 16, 0, 0, 0, time.UTC)
|
|
pool := activitypool.NewMemoryPool()
|
|
if _, err := pool.UpsertFetched(context.Background(), "provider-a", activitypool.FetchedBatch{
|
|
ObservedAt: now, ConfiguredTTL: time.Minute, MaxSize: 1,
|
|
Proxies: []proxyDomain.Proxy{{
|
|
ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080, State: proxyDomain.StateAvailable,
|
|
}},
|
|
}); err != nil {
|
|
t.Fatalf("UpsertFetched() error = %v", err)
|
|
}
|
|
assignment, err := pool.Assign(context.Background(), now, "proxy-a", "worker-a", time.Minute)
|
|
if err != nil {
|
|
t.Fatalf("Assign() error = %v", err)
|
|
}
|
|
policy := activitypool.UpstreamDrainPolicy{UpstreamID: "provider-a", Revision: 21}
|
|
reaper, err := NewConfiguredUpstreamDrainReaper(upstreamDrainPolicySourceStub{policies: []activitypool.UpstreamDrainPolicy{policy}}, pool,
|
|
UpstreamDrainReaperOptions{PollInterval: time.Second, BatchSize: 4, Now: func() time.Time { return now }})
|
|
if err != nil {
|
|
t.Fatalf("NewConfiguredUpstreamDrainReaper() error = %v", err)
|
|
}
|
|
result, err := reaper.Tick(context.Background())
|
|
if err != nil || result != (UpstreamDrainResult{Candidates: 1, Started: 1}) {
|
|
t.Fatalf("Tick() = %+v, %v", result, err)
|
|
}
|
|
if current, found, err := pool.Get(context.Background(), "proxy-a"); err != nil || !found ||
|
|
current.Epoch != assignment.Epoch || !current.Draining {
|
|
t.Fatalf("Get(after drain) = %+v, %t, %v", current, found, err)
|
|
}
|
|
}
|
|
|
|
func TestEffectiveUpstreamDrainPolicySourceUsesStableConfigOnlyRevision(t *testing.T) {
|
|
configuration := &config.Config{Upstreams: map[string]config.Upstream{
|
|
"provider-b": {Enabled: true},
|
|
"provider-a": {Enabled: false},
|
|
}}
|
|
source, err := NewEffectiveUpstreamDrainPolicySource(upstreamDrainConfigurationStub{configuration: configuration})
|
|
if err != nil {
|
|
t.Fatalf("NewEffectiveUpstreamDrainPolicySource() error = %v", err)
|
|
}
|
|
policies, err := source.ReadUpstreamDrainPolicies(context.Background())
|
|
if err != nil || len(policies) != 2 || policies[0] != (activitypool.UpstreamDrainPolicy{
|
|
UpstreamID: "provider-a", Revision: 1,
|
|
}) || policies[1] != (activitypool.UpstreamDrainPolicy{
|
|
UpstreamID: "provider-b", Revision: 1, Enabled: true,
|
|
}) {
|
|
t.Fatalf("ReadUpstreamDrainPolicies() = %+v, %v", policies, err)
|
|
}
|
|
}
|
|
|
|
func TestEffectiveUpstreamDrainPolicySourceCombinesConfigurationAndAdminState(t *testing.T) {
|
|
configuration := &config.Config{Upstreams: map[string]config.Upstream{
|
|
"provider-a": {Enabled: false},
|
|
"provider-b": {Enabled: true},
|
|
}}
|
|
source, err := NewEffectiveUpstreamDrainPolicySource(
|
|
upstreamDrainConfigurationStub{configuration: configuration, revision: 9},
|
|
upstreamDrainStateStub{snapshot: adminstate.Snapshot{
|
|
Config: &adminstate.ConfigRevision{Revision: 9},
|
|
Upstreams: []adminstate.UpstreamState{
|
|
{Name: "provider-a", Enabled: true, Revision: 7},
|
|
{Name: "provider-b", Enabled: false, Revision: 8},
|
|
},
|
|
}},
|
|
)
|
|
if err != nil {
|
|
t.Fatalf("NewEffectiveUpstreamDrainPolicySource() error = %v", err)
|
|
}
|
|
policies, err := source.ReadUpstreamDrainPolicies(context.Background())
|
|
want := []activitypool.UpstreamDrainPolicy{
|
|
{UpstreamID: "provider-a", Revision: 7},
|
|
{UpstreamID: "provider-b", Revision: 8},
|
|
}
|
|
if err != nil || len(policies) != len(want) || policies[0] != want[0] || policies[1] != want[1] {
|
|
t.Fatalf("ReadUpstreamDrainPolicies() = %+v, %v; want %+v", policies, err, want)
|
|
}
|
|
}
|