proxy-pool/internal/domain/activitypool/global_health_test.go
youfak 5ad1d10957
Some checks failed
ci / proto (push) Has been cancelled
ci / test (ubuntu-latest) (push) Has been cancelled
ci / test (windows-latest) (push) Has been cancelled
ci / race (push) Has been cancelled
ci / integration (push) Has been cancelled
feat: drain unhealthy worker assignments automatically
2026-08-02 12:38:52 +08:00

218 lines
10 KiB
Go

package activitypool
import (
"context"
"errors"
"testing"
"time"
healthDomain "proxy-pool/internal/domain/health"
proxyDomain "proxy-pool/internal/domain/proxy"
)
func TestMemoryPoolAppliesGlobalObservationAtomically(t *testing.T) {
now := time.Date(2026, 7, 31, 11, 0, 0, 0, time.UTC)
pool := seededHealthPool(t, now)
if _, err := pool.ApplyHealth(context.Background(), HealthUpdate{
ProxyID: "proxy-a", CheckedAt: now.Add(time.Second), NextState: proxyDomain.StateChecking,
}); err != nil {
t.Fatalf("ApplyHealth(checking): %v", err)
}
failed, err := pool.ApplyGlobalObservation(context.Background(), GlobalHealthCommand{
Observation: healthObservation("task-1", false, now.Add(2*time.Second)), MaxConsecutiveFailures: 3,
})
if err != nil || failed.State != proxyDomain.StateUnhealthy || failed.GlobalHealth.ConsecutiveFailures != 1 ||
!failed.GlobalHealth.UnhealthySince.Equal(now.Add(2*time.Second)) {
t.Fatalf("ApplyGlobalObservation(initial failure) = %+v, %v", failed, err)
}
if _, err := pool.ApplyHealth(context.Background(), HealthUpdate{
ProxyID: "proxy-a", CheckedAt: now.Add(3 * time.Second), NextState: proxyDomain.StateChecking,
}); err != nil {
t.Fatalf("ApplyHealth(recheck): %v", err)
}
available, err := pool.ApplyGlobalObservation(context.Background(), GlobalHealthCommand{
Observation: healthObservation("task-2", true, now.Add(4*time.Second)), MaxConsecutiveFailures: 3,
})
if err != nil || available.State != proxyDomain.StateAvailable || available.GlobalHealth.ConsecutiveFailures != 0 ||
!available.GlobalHealth.UnhealthySince.IsZero() ||
available.Proxy.LastSuccessAt == nil || !available.Proxy.LastSuccessAt.Equal(now.Add(4*time.Second)) {
t.Fatalf("ApplyGlobalObservation(recovery) = %+v, %v", available, err)
}
}
func TestMemoryPoolStartsConditionalDrainForOwnedProxyPastUnhealthyGrace(t *testing.T) {
now := time.Date(2026, 8, 2, 13, 0, 0, 0, time.UTC)
pool := seededHealthPool(t, now)
if _, err := pool.ApplyHealth(context.Background(), HealthUpdate{
ProxyID: "proxy-a", CheckedAt: now.Add(time.Second), NextState: proxyDomain.StateChecking,
}); err != nil {
t.Fatalf("ApplyHealth(checking) = %v", err)
}
if _, err := pool.ApplyGlobalObservation(context.Background(), GlobalHealthCommand{
Observation: healthObservation("task-available", true, now.Add(2*time.Second)), MaxConsecutiveFailures: 1,
}); err != nil {
t.Fatalf("ApplyGlobalObservation(available) = %v", err)
}
assignment, err := pool.Assign(context.Background(), now.Add(3*time.Second), "proxy-a", "worker-a", time.Minute)
if err != nil {
t.Fatalf("Assign() = %v", err)
}
if _, err := pool.ApplyGlobalObservation(context.Background(), GlobalHealthCommand{
Observation: healthObservation("task-suspect", false, now.Add(4*time.Second)), MaxConsecutiveFailures: 2,
}); err != nil {
t.Fatalf("ApplyGlobalObservation(suspect) = %v", err)
}
if _, err := pool.ApplyGlobalObservation(context.Background(), GlobalHealthCommand{
Observation: healthObservation("task-unhealthy", false, now.Add(5*time.Second)), MaxConsecutiveFailures: 2,
}); err != nil {
t.Fatalf("ApplyGlobalObservation(unhealthy) = %v", err)
}
command := UnhealthySweepCommand{
Now: now.Add(7 * time.Second), Limit: 8, RemoveAfterByUpstream: map[string]time.Duration{"provider-a": time.Second},
}
result, err := pool.SweepUnhealthy(context.Background(), command)
if err != nil || result.Removed != 0 || result.DeferredOwned != 1 || len(result.DrainCandidates) != 1 {
t.Fatalf("SweepUnhealthy(owned) = %+v, %v", result, err)
}
candidate := result.DrainCandidates[0]
if candidate.ProxyID != "proxy-a" || candidate.WorkerID != "worker-a" || candidate.AssignmentEpoch != assignment.Epoch ||
!candidate.UnhealthySince.Equal(now.Add(5*time.Second)) {
t.Fatalf("SweepUnhealthy() candidate = %+v", candidate)
}
if _, exists := pool.entryByIDLocked("proxy-a"); !exists {
t.Fatal("owned unhealthy proxy was removed")
}
if started, err := pool.BeginUnhealthyDrain(context.Background(), command.Now, candidate); err != nil || !started {
t.Fatalf("BeginUnhealthyDrain() = %t, %v", started, err)
}
if started, err := pool.BeginUnhealthyDrain(context.Background(), command.Now, candidate); err != nil || started {
t.Fatalf("BeginUnhealthyDrain(idempotent) = %t, %v", started, err)
}
if err := pool.AcknowledgeDrain(context.Background(), "proxy-a", "worker-a", assignment.Epoch, 0, 0); err != nil {
t.Fatalf("AcknowledgeDrain() = %v", err)
}
command.Now = now.Add(9 * time.Second)
if result, err := pool.SweepUnhealthy(context.Background(), command); err != nil || result.Removed != 1 || result.DeferredOwned != 0 || len(result.DrainCandidates) != 0 {
t.Fatalf("SweepUnhealthy(unowned) = %+v, %v", result, err)
}
if _, exists := pool.entryByIDLocked("proxy-a"); exists {
t.Fatal("unowned unhealthy proxy remains after grace")
}
}
func TestMemoryPoolRejectsUnhealthyDrainCandidateAfterRecovery(t *testing.T) {
now := time.Date(2026, 8, 2, 13, 0, 0, 0, time.UTC)
pool := seededHealthPool(t, now)
if _, err := pool.ApplyHealth(context.Background(), HealthUpdate{
ProxyID: "proxy-a", CheckedAt: now.Add(time.Second), NextState: proxyDomain.StateChecking,
}); err != nil {
t.Fatalf("ApplyHealth(checking) = %v", err)
}
if _, err := pool.ApplyGlobalObservation(context.Background(), GlobalHealthCommand{
Observation: healthObservation("task-available", true, now.Add(2*time.Second)), MaxConsecutiveFailures: 1,
}); err != nil {
t.Fatalf("ApplyGlobalObservation(available) = %v", err)
}
assignment, err := pool.Assign(context.Background(), now.Add(3*time.Second), "proxy-a", "worker-a", time.Minute)
if err != nil {
t.Fatalf("Assign() = %v", err)
}
if _, err := pool.ApplyGlobalObservation(context.Background(), GlobalHealthCommand{
Observation: healthObservation("task-suspect", false, now.Add(4*time.Second)), MaxConsecutiveFailures: 2,
}); err != nil {
t.Fatalf("ApplyGlobalObservation(suspect) = %v", err)
}
if _, err := pool.ApplyGlobalObservation(context.Background(), GlobalHealthCommand{
Observation: healthObservation("task-unhealthy", false, now.Add(5*time.Second)), MaxConsecutiveFailures: 2,
}); err != nil {
t.Fatalf("ApplyGlobalObservation(unhealthy) = %v", err)
}
command := UnhealthySweepCommand{
Now: now.Add(7 * time.Second), Limit: 1, RemoveAfterByUpstream: map[string]time.Duration{"provider-a": time.Second},
}
result, err := pool.SweepUnhealthy(context.Background(), command)
if err != nil || len(result.DrainCandidates) != 1 {
t.Fatalf("SweepUnhealthy() = %+v, %v", result, err)
}
if _, err := pool.ApplyHealth(context.Background(), HealthUpdate{
ProxyID: "proxy-a", CheckedAt: now.Add(7 * time.Second), NextState: proxyDomain.StateChecking,
}); err != nil {
t.Fatalf("ApplyHealth(rechecking) = %v", err)
}
if _, err := pool.ApplyGlobalObservation(context.Background(), GlobalHealthCommand{
Observation: healthObservation("task-recovered", true, now.Add(8*time.Second)), MaxConsecutiveFailures: 1,
}); err != nil {
t.Fatalf("ApplyGlobalObservation(recovered) = %v", err)
}
if started, err := pool.BeginUnhealthyDrain(context.Background(), now.Add(8*time.Second), result.DrainCandidates[0]); err != nil || started {
t.Fatalf("BeginUnhealthyDrain(recovered) = %t, %v", started, err)
}
if current, ok, err := pool.Get(context.Background(), "proxy-a"); err != nil || !ok || current.Epoch != assignment.Epoch || current.Draining {
t.Fatalf("Get(recovered) = %+v, %t, %v", current, ok, err)
}
}
func TestMemoryPoolGlobalObservationReplayDoesNotIncreaseFailureStreak(t *testing.T) {
now := time.Date(2026, 7, 31, 11, 0, 0, 0, time.UTC)
pool := seededHealthPool(t, now)
if _, err := pool.ApplyHealth(context.Background(), HealthUpdate{
ProxyID: "proxy-a", CheckedAt: now.Add(time.Second), NextState: proxyDomain.StateChecking,
}); err != nil {
t.Fatalf("ApplyHealth(checking): %v", err)
}
command := GlobalHealthCommand{Observation: healthObservation("task-1", false, now.Add(2*time.Second)), MaxConsecutiveFailures: 3}
first, err := pool.ApplyGlobalObservation(context.Background(), command)
if err != nil {
t.Fatalf("ApplyGlobalObservation(first): %v", err)
}
replayed, err := pool.ApplyGlobalObservation(context.Background(), command)
if err != nil || replayed.GlobalHealth != first.GlobalHealth {
t.Fatalf("ApplyGlobalObservation(replay) = %+v, %v", replayed, err)
}
conflicting := command
conflicting.Observation.Success = true
conflicting.Observation.FailureClass = ""
if _, err := pool.ApplyGlobalObservation(context.Background(), conflicting); !errors.Is(err, healthDomain.ErrConflictingObservation) {
t.Fatalf("ApplyGlobalObservation(conflict) error = %v, want ErrConflictingObservation", err)
}
}
func TestMemoryPoolRejectsTargetObservationFromGlobalStore(t *testing.T) {
now := time.Date(2026, 7, 31, 11, 0, 0, 0, time.UTC)
pool := seededHealthPool(t, now)
_, err := pool.ApplyGlobalObservation(context.Background(), GlobalHealthCommand{Observation: healthDomain.Observation{
TaskID: "task-target", ProxyID: "proxy-a", Level: healthDomain.LevelTarget, RoutingName: "route-a",
TargetURL: "https://target.example/", Success: false, FailureClass: "target_403", ObservedAt: now.Add(time.Second),
}, MaxConsecutiveFailures: 2})
if !errors.Is(err, ErrInvalidHealthUpdate) {
t.Fatalf("ApplyGlobalObservation(target) error = %v, want ErrInvalidHealthUpdate", err)
}
}
func seededHealthPool(t *testing.T, now time.Time) *MemoryPool {
t.Helper()
pool := NewMemoryPool()
_, err := pool.UpsertFetched(context.Background(), "provider-a", FetchedBatch{
ObservedAt: now, ConfiguredTTL: time.Minute, MaxSize: 1,
Proxies: []proxyDomain.Proxy{{
ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080,
}},
})
if err != nil {
t.Fatalf("UpsertFetched(): %v", err)
}
return pool
}
func healthObservation(taskID string, success bool, observedAt time.Time) healthDomain.Observation {
observation := healthDomain.Observation{
TaskID: taskID, ProxyID: "proxy-a", Level: healthDomain.LevelBasic,
Success: success, Latency: 10 * time.Millisecond, ObservedAt: observedAt,
}
if !success {
observation.FailureClass = "timeout"
}
return observation
}