264 lines
11 KiB
Go
264 lines
11 KiB
Go
package health
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
|
|
"proxy-pool/internal/config"
|
|
healthDomain "proxy-pool/internal/domain/health"
|
|
proxyDomain "proxy-pool/internal/domain/proxy"
|
|
)
|
|
|
|
func TestSchedulerRunnerUsesBoundedDueBatchAndSharedInFlightCount(t *testing.T) {
|
|
planner, err := NewPlanner(SchedulePolicy{
|
|
Interval: time.Minute, MaxInFlight: 3, Timeout: time.Second, MaxAttempts: 1,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewPlanner(): %v", err)
|
|
}
|
|
now := time.Date(2026, 7, 31, 18, 0, 0, 0, time.UTC)
|
|
source := &dueSourceStub{inFlight: 1, candidates: []Candidate{
|
|
{ProxyID: "fetched", State: proxyDomain.StateFetched, Level: healthDomain.LevelBasic, DueAt: now},
|
|
{ProxyID: "available", State: proxyDomain.StateAvailable, Level: healthDomain.LevelBasic, DueAt: now},
|
|
}}
|
|
sink := &taskSinkStub{}
|
|
runner, err := NewSchedulerRunner(planner, source, sink, SchedulerRunnerOptions{
|
|
PollInterval: time.Second, BatchSize: 2, Now: func() time.Time { return now },
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewSchedulerRunner(): %v", err)
|
|
}
|
|
result, err := runner.Tick(context.Background())
|
|
if err != nil || result.Planned != 2 || result.Offered != 2 || source.limit != 2 || len(sink.tasks) != 2 ||
|
|
sink.tasks[0].Priority != PriorityFetched {
|
|
t.Fatalf("Tick() = (%+v, %v); source=%+v sink=%+v", result, err, source, sink.tasks)
|
|
}
|
|
}
|
|
|
|
func TestSchedulerRunnerSkipsDueReadAtSharedCapacity(t *testing.T) {
|
|
planner, err := NewPlanner(SchedulePolicy{
|
|
Interval: time.Minute, MaxInFlight: 1, Timeout: time.Second, MaxAttempts: 1,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewPlanner(): %v", err)
|
|
}
|
|
source := &dueSourceStub{inFlight: 1}
|
|
runner, err := NewSchedulerRunner(planner, source, &taskSinkStub{}, SchedulerRunnerOptions{
|
|
PollInterval: time.Second, BatchSize: 1, Now: time.Now,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewSchedulerRunner(): %v", err)
|
|
}
|
|
result, err := runner.Tick(context.Background())
|
|
if err != nil || result != (TickResult{}) || source.limit != 0 {
|
|
t.Fatalf("Tick(at capacity) = (%+v, %v); source=%+v", result, err, source)
|
|
}
|
|
}
|
|
|
|
func TestSchedulerRunnerRejectsOversizedSourceAndSinkResponses(t *testing.T) {
|
|
planner, err := NewPlanner(SchedulePolicy{
|
|
Interval: time.Minute, MaxInFlight: 2, Timeout: time.Second, MaxAttempts: 1,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewPlanner(): %v", err)
|
|
}
|
|
now := time.Now()
|
|
runner, err := NewSchedulerRunner(planner, &dueSourceStub{candidates: []Candidate{
|
|
{ProxyID: "one", State: proxyDomain.StateFetched, Level: healthDomain.LevelBasic, DueAt: now},
|
|
{ProxyID: "two", State: proxyDomain.StateFetched, Level: healthDomain.LevelBasic, DueAt: now},
|
|
}}, &taskSinkStub{}, SchedulerRunnerOptions{PollInterval: time.Second, BatchSize: 1, Now: func() time.Time { return now }})
|
|
if err != nil {
|
|
t.Fatalf("NewSchedulerRunner(): %v", err)
|
|
}
|
|
if _, err := runner.Tick(context.Background()); !errors.Is(err, ErrInvalidDueSource) {
|
|
t.Fatalf("Tick(oversized source) error = %v, want ErrInvalidDueSource", err)
|
|
}
|
|
}
|
|
|
|
func TestConfiguredSchedulerRunnerFollowsUpstreamConfigChanges(t *testing.T) {
|
|
now := time.Date(2026, 8, 2, 10, 0, 0, 0, time.UTC)
|
|
configuration := &config.Config{
|
|
Defaults: config.Defaults{Check: config.Check{
|
|
Interval: config.Duration(time.Minute),
|
|
MaxInFlight: 2,
|
|
Timeout: config.Duration(time.Second),
|
|
MaxAttempts: 1,
|
|
}},
|
|
Upstreams: map[string]config.Upstream{
|
|
"provider-a": {Enabled: true},
|
|
},
|
|
}
|
|
source := &dueSourceStub{candidates: []Candidate{{
|
|
ProxyID: "proxy-a", State: proxyDomain.StateFetched, Level: healthDomain.LevelBasic, DueAt: now,
|
|
}}}
|
|
sink := &taskSinkStub{}
|
|
runner, err := NewConfiguredSchedulerRunner(&configurationSourceStub{configuration: configuration}, "provider-a", source, sink,
|
|
SchedulerRunnerOptions{PollInterval: time.Second, BatchSize: 16, Now: func() time.Time { return now }})
|
|
if err != nil {
|
|
t.Fatalf("NewConfiguredSchedulerRunner(): %v", err)
|
|
}
|
|
|
|
result, err := runner.Tick(context.Background())
|
|
if err != nil || result != (TickResult{Planned: 1, Offered: 1}) || len(sink.tasks) != 1 || sink.tasks[0].Attempts != 1 ||
|
|
source.limit != 2 {
|
|
t.Fatalf("initial Tick() = (%+v, %v); source=%+v sink=%+v", result, err, source, sink.tasks)
|
|
}
|
|
|
|
updated := configuration.Upstreams["provider-a"]
|
|
updated.Check.MaxAttempts = 3
|
|
configuration.Upstreams["provider-a"] = updated
|
|
result, err = runner.Tick(context.Background())
|
|
if err != nil || result != (TickResult{Planned: 1, Offered: 1}) || len(sink.tasks) != 1 || sink.tasks[0].Attempts != 3 {
|
|
t.Fatalf("updated Tick() = (%+v, %v); sink=%+v", result, err, sink.tasks)
|
|
}
|
|
|
|
updated.Enabled = false
|
|
configuration.Upstreams["provider-a"] = updated
|
|
inFlightCalls, dueCalls := source.inFlightCalls, source.dueCalls
|
|
result, err = runner.Tick(context.Background())
|
|
if err != nil || result != (TickResult{}) || source.inFlightCalls != inFlightCalls || source.dueCalls != dueCalls {
|
|
t.Fatalf("disabled Tick() = (%+v, %v); source=%+v", result, err, source)
|
|
}
|
|
}
|
|
|
|
func TestConfiguredSchedulerSupervisorDiscoversNewEnabledUpstream(t *testing.T) {
|
|
now := time.Date(2026, 8, 2, 10, 0, 0, 0, time.UTC)
|
|
configuration := &config.Config{
|
|
Defaults: config.Defaults{Check: config.Check{
|
|
Interval: config.Duration(time.Minute),
|
|
MaxInFlight: 2,
|
|
Timeout: config.Duration(time.Second),
|
|
MaxAttempts: 1,
|
|
}},
|
|
Upstreams: map[string]config.Upstream{
|
|
"provider-a": {Enabled: true},
|
|
},
|
|
}
|
|
source := &upstreamTaskSourceStub{sources: map[string]*dueSourceStub{
|
|
"provider-a": {candidates: []Candidate{{
|
|
ProxyID: "proxy-a", State: proxyDomain.StateFetched, Level: healthDomain.LevelBasic, DueAt: now,
|
|
}}},
|
|
"provider-b": {candidates: []Candidate{{
|
|
ProxyID: "proxy-b", State: proxyDomain.StateFetched, Level: healthDomain.LevelBasic, DueAt: now,
|
|
}}},
|
|
}}
|
|
sink := &taskSinkStub{}
|
|
supervisor, err := NewConfiguredSchedulerSupervisor(&configurationSourceStub{configuration: configuration}, source, sink,
|
|
SchedulerRunnerOptions{PollInterval: time.Second, BatchSize: 16, Now: func() time.Time { return now }})
|
|
if err != nil {
|
|
t.Fatalf("NewConfiguredSchedulerSupervisor(): %v", err)
|
|
}
|
|
if result, err := supervisor.Tick(context.Background()); err != nil || result != (TickResult{Planned: 1, Offered: 1}) ||
|
|
source.sources["provider-a"].dueCalls != 1 || source.sources["provider-b"].dueCalls != 0 {
|
|
t.Fatalf("initial Tick() = (%+v, %v); source=%+v", result, err, source)
|
|
}
|
|
|
|
configuration.Upstreams["provider-a"] = config.Upstream{Enabled: false}
|
|
configuration.Upstreams["provider-b"] = config.Upstream{Enabled: true, Check: config.Check{MaxAttempts: 3, MaxInFlight: 1}}
|
|
result, err := supervisor.Tick(context.Background())
|
|
if err != nil || result != (TickResult{Planned: 1, Offered: 1}) || source.sources["provider-b"].dueCalls != 1 ||
|
|
len(sink.tasks) != 1 || sink.tasks[0].Candidate.ProxyID != "proxy-b" || sink.tasks[0].Attempts != 3 || source.sources["provider-b"].limit != 1 {
|
|
t.Fatalf("updated Tick() = (%+v, %v); source=%+v sink=%+v", result, err, source, sink.tasks)
|
|
}
|
|
}
|
|
|
|
func TestConfiguredSchedulerSupervisorSchedulesBoundedEgressGroups(t *testing.T) {
|
|
now := time.Date(2026, 8, 2, 11, 0, 0, 0, time.UTC)
|
|
configuration := &config.Config{
|
|
Defaults: config.Defaults{Check: config.Check{
|
|
Interval: config.Duration(time.Minute), MaxInFlight: 6, Timeout: config.Duration(time.Second), MaxAttempts: 1,
|
|
URLs: []string{"https://egress-one.example/identity", "https://egress-two.example/identity"},
|
|
}},
|
|
Upstreams: map[string]config.Upstream{"provider-a": {Enabled: true}},
|
|
}
|
|
source := &upstreamTaskSourceStub{sources: map[string]*dueSourceStub{
|
|
"provider-a": {candidates: []Candidate{{
|
|
ProxyID: "basic", State: proxyDomain.StateFetched, Level: healthDomain.LevelBasic, DueAt: now,
|
|
}}, egressCandidates: map[string][]Candidate{
|
|
"https://egress-one.example/identity": {{ProxyID: "egress-one", State: proxyDomain.StateFetched, Level: healthDomain.LevelEgress, TargetURL: "https://egress-one.example/identity", DueAt: now}},
|
|
"https://egress-two.example/identity": {{ProxyID: "egress-two", State: proxyDomain.StateFetched, Level: healthDomain.LevelEgress, TargetURL: "https://egress-two.example/identity", DueAt: now}},
|
|
}},
|
|
}}
|
|
sink := &taskSinkStub{}
|
|
supervisor, err := NewConfiguredSchedulerSupervisor(&configurationSourceStub{configuration: configuration}, source, sink,
|
|
SchedulerRunnerOptions{PollInterval: time.Second, BatchSize: 6, Now: func() time.Time { return now }})
|
|
if err != nil {
|
|
t.Fatalf("NewConfiguredSchedulerSupervisor() = %v", err)
|
|
}
|
|
result, err := supervisor.Tick(context.Background())
|
|
if err != nil || result != (TickResult{Planned: 3, Offered: 3}) || len(sink.tasks) != 1 ||
|
|
sink.tasks[0].Candidate.Level != healthDomain.LevelEgress || source.sources["provider-a"].egressCalls != 2 {
|
|
t.Fatalf("Tick() = (%+v, %v); source=%+v sink=%+v", result, err, source, sink.tasks)
|
|
}
|
|
}
|
|
|
|
type dueSourceStub struct {
|
|
inFlight int
|
|
candidates []Candidate
|
|
limit int
|
|
inFlightCalls int
|
|
dueCalls int
|
|
egressCandidates map[string][]Candidate
|
|
egressCalls int
|
|
}
|
|
|
|
func (source *dueSourceStub) InFlight(context.Context, time.Time) (int, error) {
|
|
source.inFlightCalls++
|
|
return source.inFlight, nil
|
|
}
|
|
|
|
func (source *dueSourceStub) DueCandidates(_ context.Context, _ time.Time, limit int) ([]Candidate, error) {
|
|
source.dueCalls++
|
|
source.limit = limit
|
|
return source.candidates, nil
|
|
}
|
|
|
|
func (source *dueSourceStub) DueEgressCandidates(_ context.Context, targetURL string, _ time.Time, limit int) ([]Candidate, error) {
|
|
source.egressCalls++
|
|
source.limit = limit
|
|
return source.egressCandidates[targetURL], nil
|
|
}
|
|
|
|
type upstreamTaskSourceStub struct {
|
|
sources map[string]*dueSourceStub
|
|
}
|
|
|
|
func (source *upstreamTaskSourceStub) InFlightForUpstream(ctx context.Context, upstreamID string, now time.Time) (int, error) {
|
|
item, exists := source.sources[upstreamID]
|
|
if !exists {
|
|
return 0, errors.New("missing upstream source")
|
|
}
|
|
return item.InFlight(ctx, now)
|
|
}
|
|
|
|
func (source *upstreamTaskSourceStub) DueCandidatesForUpstream(ctx context.Context, upstreamID string, now time.Time, limit int) ([]Candidate, error) {
|
|
item, exists := source.sources[upstreamID]
|
|
if !exists {
|
|
return nil, errors.New("missing upstream source")
|
|
}
|
|
return item.DueCandidates(ctx, now, limit)
|
|
}
|
|
|
|
func (source *upstreamTaskSourceStub) DueEgressCandidatesForUpstream(ctx context.Context, upstreamID, targetURL string, now time.Time, limit int) ([]Candidate, error) {
|
|
item, exists := source.sources[upstreamID]
|
|
if !exists {
|
|
return nil, errors.New("missing upstream source")
|
|
}
|
|
return item.DueEgressCandidates(ctx, targetURL, now, limit)
|
|
}
|
|
|
|
type taskSinkStub struct {
|
|
tasks []PlannedTask
|
|
offer int
|
|
}
|
|
|
|
func (sink *taskSinkStub) Offer(_ context.Context, tasks []PlannedTask) (int, error) {
|
|
sink.tasks = append([]PlannedTask(nil), tasks...)
|
|
if sink.offer != 0 {
|
|
return sink.offer, nil
|
|
}
|
|
return len(tasks), nil
|
|
}
|