proxy-pool/internal/adapters/redisactivity/global_health_integration_test.go

//go:build integration

package redisactivity

import (
	"context"
	"errors"
	"testing"
	"time"

	"github.com/redis/go-redis/v9"

	"proxy-pool/internal/domain/activitypool"
	healthDomain "proxy-pool/internal/domain/health"
	proxyDomain "proxy-pool/internal/domain/proxy"
)

func TestRedisAppliesGlobalHealthObservationAtomically(t *testing.T) {
	fixture := newRedisTestFixture(t)
	now := redisTestNow()
	if _, err := fixture.Adapter.UpsertFetched(context.Background(), "provider-a", activitypool.FetchedBatch{
		ObservedAt: now, ConfiguredTTL: time.Minute, MaxSize: 1,
		Proxies: []proxyDomain.Proxy{testProxy("proxy-a", "192.0.2.10")},
	}); err != nil {
		t.Fatalf("UpsertFetched(): %v", err)
	}
	if _, err := fixture.Adapter.ApplyHealth(context.Background(), activitypool.HealthUpdate{
		ProxyID: "proxy-a", CheckedAt: now.Add(time.Second), NextState: proxyDomain.StateChecking,
	}); err != nil {
		t.Fatalf("ApplyHealth(checking): %v", err)
	}
	first := redisGlobalObservation("task-1", false, now.Add(2*time.Second))
	failed, err := fixture.Adapter.ApplyGlobalObservation(context.Background(), activitypool.GlobalHealthCommand{
		Observation: first, 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)
	}
	replayed, err := fixture.Adapter.ApplyGlobalObservation(context.Background(), activitypool.GlobalHealthCommand{
		Observation: first, MaxConsecutiveFailures: 3,
	})
	if err != nil || replayed.GlobalHealth != failed.GlobalHealth {
		t.Fatalf("ApplyGlobalObservation(replay) = %+v, %v", replayed, err)
	}
	conflicting := first
	conflicting.Success = true
	conflicting.FailureClass = ""
	if _, err := fixture.Adapter.ApplyGlobalObservation(context.Background(), activitypool.GlobalHealthCommand{
		Observation: conflicting, MaxConsecutiveFailures: 3,
	}); !errors.Is(err, healthDomain.ErrConflictingObservation) {
		t.Fatalf("ApplyGlobalObservation(conflict) error = %v, want ErrConflictingObservation", err)
	}
	if _, err := fixture.Adapter.ApplyHealth(context.Background(), activitypool.HealthUpdate{
		ProxyID: "proxy-a", CheckedAt: now.Add(3 * time.Second), NextState: proxyDomain.StateChecking,
	}); err != nil {
		t.Fatalf("ApplyHealth(recheck): %v", err)
	}
	recovered, err := fixture.Adapter.ApplyGlobalObservation(context.Background(), activitypool.GlobalHealthCommand{
		Observation: redisGlobalObservation("task-2", true, now.Add(4*time.Second)), MaxConsecutiveFailures: 3,
	})
	if err != nil || recovered.State != proxyDomain.StateAvailable || recovered.GlobalHealth.ConsecutiveFailures != 0 ||
		!recovered.GlobalHealth.UnhealthySince.IsZero() ||
		recovered.GlobalHealth.LastObservedAt != now.Add(4*time.Second) {
		t.Fatalf("ApplyGlobalObservation(recovery) = %+v, %v", recovered, err)
	}
}

func TestRedisSweepsUnownedProxiesAfterUnhealthyGrace(t *testing.T) {
	fixture := newRedisTestFixture(t)
	now := redisTestNow()
	if _, err := fixture.Adapter.UpsertFetched(context.Background(), "provider-a", activitypool.FetchedBatch{
		ObservedAt: now, ConfiguredTTL: time.Minute, MaxSize: 1,
		Proxies: []proxyDomain.Proxy{testProxy("proxy-a", "192.0.2.10")},
	}); err != nil {
		t.Fatalf("UpsertFetched() = %v", err)
	}
	if _, err := fixture.Adapter.ApplyHealth(context.Background(), activitypool.HealthUpdate{
		ProxyID: "proxy-a", CheckedAt: now.Add(time.Second), NextState: proxyDomain.StateChecking,
	}); err != nil {
		t.Fatalf("ApplyHealth(checking) = %v", err)
	}
	if _, err := fixture.Adapter.ApplyGlobalObservation(context.Background(), activitypool.GlobalHealthCommand{
		Observation: redisGlobalObservation("task-unhealthy", false, now.Add(2*time.Second)), MaxConsecutiveFailures: 2,
	}); err != nil {
		t.Fatalf("ApplyGlobalObservation(unhealthy) = %v", err)
	}
	if _, err := fixture.Client.ZScore(context.Background(), fixture.Adapter.keys.healthUnhealthy, "proxy-a").Result(); err != nil {
		t.Fatalf("health unhealthy index = %v", err)
	}
	command := activitypool.UnhealthySweepCommand{
		Now: now.Add(3 * time.Second), Limit: 4, RemoveAfterByUpstream: map[string]time.Duration{"provider-a": 3 * time.Second},
	}
	if result, err := fixture.Adapter.SweepUnhealthy(context.Background(), command); err != nil || result.Removed != 0 || result.DeferredOwned != 0 || len(result.DrainCandidates) != 0 {
		t.Fatalf("SweepUnhealthy(before grace) = %+v, %v", result, err)
	}
	command.Now = now.Add(6 * time.Second)
	if result, err := fixture.Adapter.SweepUnhealthy(context.Background(), command); err != nil || result.Removed != 1 || result.DeferredOwned != 0 || len(result.DrainCandidates) != 0 {
		t.Fatalf("SweepUnhealthy(after grace) = %+v, %v", result, err)
	}
	if _, err := fixture.Adapter.UpstreamForProxy(context.Background(), "proxy-a", command.Now); !errors.Is(err, activitypool.ErrActivityNotFound) {
		t.Fatalf("UpstreamForProxy(removed) error = %v, want ErrActivityNotFound", err)
	}
	if _, err := fixture.Client.ZScore(context.Background(), fixture.Adapter.keys.healthUnhealthy, "proxy-a").Result(); !errors.Is(err, redis.Nil) {
		t.Fatalf("health unhealthy index after removal = %v, want redis.Nil", err)
	}
}

func TestRedisUnhealthySweepStartsConditionalDrainForOwnedProxy(t *testing.T) {
	fixture := newRedisTestFixture(t)
	now := redisTestNow()
	if _, err := fixture.Adapter.UpsertFetched(context.Background(), "provider-a", activitypool.FetchedBatch{
		ObservedAt: now, ConfiguredTTL: time.Minute, MaxSize: 1,
		Proxies: []proxyDomain.Proxy{testProxy("proxy-a", "192.0.2.10")},
	}); err != nil {
		t.Fatalf("UpsertFetched() = %v", err)
	}
	if _, err := fixture.Adapter.ApplyHealth(context.Background(), activitypool.HealthUpdate{
		ProxyID: "proxy-a", CheckedAt: now.Add(time.Second), NextState: proxyDomain.StateChecking,
	}); err != nil {
		t.Fatalf("ApplyHealth(checking) = %v", err)
	}
	if _, err := fixture.Adapter.ApplyGlobalObservation(context.Background(), activitypool.GlobalHealthCommand{
		Observation: redisGlobalObservation("task-available", true, now.Add(2*time.Second)), MaxConsecutiveFailures: 2,
	}); err != nil {
		t.Fatalf("ApplyGlobalObservation(available) = %v", err)
	}
	assignment, err := fixture.Adapter.Assign(context.Background(), now.Add(3*time.Second), "proxy-a", "worker-a", time.Minute)
	if err != nil {
		t.Fatalf("Assign() = %v", err)
	}
	if _, err := fixture.Adapter.ApplyGlobalObservation(context.Background(), activitypool.GlobalHealthCommand{
		Observation: redisGlobalObservation("task-suspect", false, now.Add(4*time.Second)), MaxConsecutiveFailures: 2,
	}); err != nil {
		t.Fatalf("ApplyGlobalObservation(suspect) = %v", err)
	}
	if _, err := fixture.Adapter.ApplyGlobalObservation(context.Background(), activitypool.GlobalHealthCommand{
		Observation: redisGlobalObservation("task-unhealthy", false, now.Add(5*time.Second)), MaxConsecutiveFailures: 2,
	}); err != nil {
		t.Fatalf("ApplyGlobalObservation(unhealthy) = %v", err)
	}
	command := activitypool.UnhealthySweepCommand{
		Now: now.Add(7 * time.Second), Limit: 4, RemoveAfterByUpstream: map[string]time.Duration{"provider-a": time.Second},
	}
	result, err := fixture.Adapter.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 started, err := fixture.Adapter.BeginUnhealthyDrain(context.Background(), command.Now, candidate); err != nil || !started {
		t.Fatalf("BeginUnhealthyDrain() = %t, %v", started, err)
	}
	if started, err := fixture.Adapter.BeginUnhealthyDrain(context.Background(), command.Now, candidate); err != nil || started {
		t.Fatalf("BeginUnhealthyDrain(idempotent) = %t, %v", started, err)
	}
	if err := fixture.Adapter.AcknowledgeDrain(context.Background(), "proxy-a", "worker-a", assignment.Epoch, 0, 0); err != nil {
		t.Fatalf("AcknowledgeDrain() = %v", err)
	}
	command.Now = now.Add(8 * time.Second)
	if result, err := fixture.Adapter.SweepUnhealthy(context.Background(), command); err != nil || result.Removed != 1 || result.DeferredOwned != 0 || len(result.DrainCandidates) != 0 {
		t.Fatalf("SweepUnhealthy(unowned) = %+v, %v", result, err)
	}
}

func TestRedisRejectsUnhealthyDrainCandidateAfterRecovery(t *testing.T) {
	fixture := newRedisTestFixture(t)
	now := redisTestNow()
	if _, err := fixture.Adapter.UpsertFetched(context.Background(), "provider-a", activitypool.FetchedBatch{
		ObservedAt: now, ConfiguredTTL: time.Minute, MaxSize: 1,
		Proxies: []proxyDomain.Proxy{testProxy("proxy-a", "192.0.2.10")},
	}); err != nil {
		t.Fatalf("UpsertFetched() = %v", err)
	}
	if _, err := fixture.Adapter.ApplyHealth(context.Background(), activitypool.HealthUpdate{
		ProxyID: "proxy-a", CheckedAt: now.Add(time.Second), NextState: proxyDomain.StateChecking,
	}); err != nil {
		t.Fatalf("ApplyHealth(checking) = %v", err)
	}
	if _, err := fixture.Adapter.ApplyGlobalObservation(context.Background(), activitypool.GlobalHealthCommand{
		Observation: redisGlobalObservation("task-available", true, now.Add(2*time.Second)), MaxConsecutiveFailures: 2,
	}); err != nil {
		t.Fatalf("ApplyGlobalObservation(available) = %v", err)
	}
	assignment, err := fixture.Adapter.Assign(context.Background(), now.Add(3*time.Second), "proxy-a", "worker-a", time.Minute)
	if err != nil {
		t.Fatalf("Assign() = %v", err)
	}
	if _, err := fixture.Adapter.ApplyGlobalObservation(context.Background(), activitypool.GlobalHealthCommand{
		Observation: redisGlobalObservation("task-suspect", false, now.Add(4*time.Second)), MaxConsecutiveFailures: 2,
	}); err != nil {
		t.Fatalf("ApplyGlobalObservation(suspect) = %v", err)
	}
	if _, err := fixture.Adapter.ApplyGlobalObservation(context.Background(), activitypool.GlobalHealthCommand{
		Observation: redisGlobalObservation("task-unhealthy", false, now.Add(5*time.Second)), MaxConsecutiveFailures: 2,
	}); err != nil {
		t.Fatalf("ApplyGlobalObservation(unhealthy) = %v", err)
	}
	command := activitypool.UnhealthySweepCommand{
		Now: now.Add(7 * time.Second), Limit: 1, RemoveAfterByUpstream: map[string]time.Duration{"provider-a": time.Second},
	}
	result, err := fixture.Adapter.SweepUnhealthy(context.Background(), command)
	if err != nil || len(result.DrainCandidates) != 1 {
		t.Fatalf("SweepUnhealthy() = %+v, %v", result, err)
	}
	if _, err := fixture.Adapter.ApplyHealth(context.Background(), activitypool.HealthUpdate{
		ProxyID: "proxy-a", CheckedAt: now.Add(8 * time.Second), NextState: proxyDomain.StateChecking,
	}); err != nil {
		t.Fatalf("ApplyHealth(rechecking) = %v", err)
	}
	if _, err := fixture.Adapter.ApplyGlobalObservation(context.Background(), activitypool.GlobalHealthCommand{
		Observation: redisGlobalObservation("task-recovered", true, now.Add(9*time.Second)), MaxConsecutiveFailures: 2,
	}); err != nil {
		t.Fatalf("ApplyGlobalObservation(recovered) = %v", err)
	}
	if started, err := fixture.Adapter.BeginUnhealthyDrain(context.Background(), now.Add(9*time.Second), result.DrainCandidates[0]); err != nil || started {
		t.Fatalf("BeginUnhealthyDrain(recovered) = %t, %v", started, err)
	}
	if current, ok, err := fixture.Adapter.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 TestRedisDisabledUpstreamDrainFencesPolicyRevision(t *testing.T) {
	fixture := newRedisTestFixture(t)
	now := redisTestNow()
	seedRedisAvailable(t, fixture.Adapter, "provider-a", now, now.Add(time.Second), time.Minute,
		testProxy("proxy-a", "192.0.2.10"))
	assignment, err := fixture.Adapter.Assign(context.Background(), now.Add(2*time.Second), "proxy-a", "worker-a", time.Minute)
	if err != nil {
		t.Fatalf("Assign() error = %v", err)
	}
	disabled := activitypool.UpstreamDrainPolicy{UpstreamID: "provider-a", Revision: 31}
	if err := fixture.Adapter.ReplaceUpstreamDrainPolicies(context.Background(), []activitypool.UpstreamDrainPolicy{disabled}); err != nil {
		t.Fatalf("ReplaceUpstreamDrainPolicies(disabled) error = %v", err)
	}
	candidates, err := fixture.Adapter.ListDisabledUpstreamDrainCandidates(context.Background(), now.Add(3*time.Second), disabled, 1)
	if err != nil || len(candidates) != 1 || candidates[0] != (activitypool.DisabledUpstreamDrainCandidate{
		UpstreamID: "provider-a", PolicyRevision: 31, ProxyID: "proxy-a", WorkerID: "worker-a", AssignmentEpoch: assignment.Epoch,
	}) {
		t.Fatalf("ListDisabledUpstreamDrainCandidates() = %+v, %v", candidates, err)
	}
	if err := fixture.Adapter.ReplaceUpstreamDrainPolicies(context.Background(), []activitypool.UpstreamDrainPolicy{{
		UpstreamID: "provider-a", Revision: 32, Enabled: true,
	}}); err != nil {
		t.Fatalf("ReplaceUpstreamDrainPolicies(enabled) error = %v", err)
	}
	if started, err := fixture.Adapter.BeginDisabledUpstreamDrain(context.Background(), now.Add(3*time.Second), candidates[0]); err != nil || started {
		t.Fatalf("BeginDisabledUpstreamDrain(stale policy) = %t, %v", started, err)
	}
	if current, ok, err := fixture.Adapter.Get(context.Background(), "proxy-a"); err != nil || !ok || current.Draining {
		t.Fatalf("Get(after stale candidate) = %+v, %t, %v", current, ok, err)
	}

	disabled.Revision = 33
	if err := fixture.Adapter.ReplaceUpstreamDrainPolicies(context.Background(), []activitypool.UpstreamDrainPolicy{disabled}); err != nil {
		t.Fatalf("ReplaceUpstreamDrainPolicies(disabled again) error = %v", err)
	}
	candidates, err = fixture.Adapter.ListDisabledUpstreamDrainCandidates(context.Background(), now.Add(4*time.Second), disabled, 1)
	if err != nil || len(candidates) != 1 {
		t.Fatalf("ListDisabledUpstreamDrainCandidates(disabled again) = %+v, %v", candidates, err)
	}
	if started, err := fixture.Adapter.BeginDisabledUpstreamDrain(context.Background(), now.Add(4*time.Second), candidates[0]); err != nil || !started {
		t.Fatalf("BeginDisabledUpstreamDrain() = %t, %v", started, err)
	}
	if current, ok, err := fixture.Adapter.Get(context.Background(), "proxy-a"); err != nil || !ok || !current.Draining {
		t.Fatalf("Get(after drain) = %+v, %t, %v", current, ok, err)
	}
}

func redisGlobalObservation(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
}