109 lines
5.4 KiB
Go
109 lines
5.4 KiB
Go
package health
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
|
|
healthDomain "proxy-pool/internal/domain/health"
|
|
proxyDomain "proxy-pool/internal/domain/proxy"
|
|
)
|
|
|
|
func TestMemoryTaskBrokerClaimsSupportedTasksAndFencesObservations(t *testing.T) {
|
|
now := time.Date(2026, 8, 1, 9, 0, 0, 0, time.UTC)
|
|
broker, err := NewMemoryTaskBroker(MemoryTaskBrokerOptions{
|
|
LeaseTTL: time.Minute,
|
|
Now: func() time.Time { return now },
|
|
Material: TaskMaterialResolverFunc(func(_ context.Context, candidate Candidate) (TaskMaterial, error) {
|
|
return TaskMaterial{
|
|
Protocol: proxyDomain.SchemeHTTP, Host: candidate.ProxyID + ".example", Port: 8080,
|
|
SecretRef: "cred-" + candidate.ProxyID, CredentialVersion: "v1", Username: "user", Password: "secret",
|
|
}, nil
|
|
}),
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewMemoryTaskBroker(): %v", err)
|
|
}
|
|
plans := []PlannedTask{
|
|
{Candidate: Candidate{ProxyID: "basic", State: proxyDomain.StateFetched, Level: healthDomain.LevelBasic, DueAt: now}, Deadline: now.Add(10 * time.Second), Attempts: 2},
|
|
{Candidate: Candidate{ProxyID: "target", State: proxyDomain.StateAvailable, Level: healthDomain.LevelTarget, RoutingName: "route-a", TargetURL: "https://target.example/check", DueAt: now}, Deadline: now.Add(10 * time.Second), Attempts: 2},
|
|
}
|
|
if offered, err := broker.Offer(context.Background(), plans); err != nil || offered != 2 {
|
|
t.Fatalf("Offer() = (%d, %v)", offered, err)
|
|
}
|
|
claimed, err := broker.Claim(context.Background(), TaskClaim{
|
|
CheckerID: "checker-a", InstanceID: "instance-a", MaxInFlight: 2, SupportedLevels: []healthDomain.Level{healthDomain.LevelBasic},
|
|
})
|
|
if err != nil || len(claimed) != 1 || claimed[0].ProxyID != "basic" || claimed[0].Password != "secret" {
|
|
t.Fatalf("Claim() = (%+v, %v)", claimed, err)
|
|
}
|
|
observation := healthDomain.Observation{
|
|
TaskID: claimed[0].TaskID, ProxyID: "basic", Level: healthDomain.LevelBasic,
|
|
Success: true, Latency: time.Millisecond, ObservedAt: now.Add(time.Second),
|
|
}
|
|
if err := broker.AuthorizeObservation(context.Background(), "checker-b", claimed[0].LeaseToken, observation, now.Add(time.Second)); !errors.Is(err, ErrTaskLeaseNotOwned) {
|
|
t.Fatalf("AuthorizeObservation(other checker) error = %v, want ErrTaskLeaseNotOwned", err)
|
|
}
|
|
if err := broker.AuthorizeObservation(context.Background(), "checker-a", claimed[0].LeaseToken, observation, now.Add(time.Second)); err != nil {
|
|
t.Fatalf("AuthorizeObservation(owner): %v", err)
|
|
}
|
|
if err := broker.CompleteObservation(context.Background(), "checker-a", claimed[0].LeaseToken, observation, now.Add(time.Second)); err != nil {
|
|
t.Fatalf("CompleteObservation(): %v", err)
|
|
}
|
|
if err := broker.AuthorizeObservation(context.Background(), "checker-a", claimed[0].LeaseToken, observation, now.Add(2*time.Second)); err != nil {
|
|
t.Fatalf("AuthorizeObservation(replay): %v", err)
|
|
}
|
|
}
|
|
|
|
func TestMemoryTaskBrokerReclaimsExpiredLeaseWithoutExceedingCheckerLimit(t *testing.T) {
|
|
now := time.Date(2026, 8, 1, 10, 0, 0, 0, time.UTC)
|
|
broker, err := NewMemoryTaskBroker(MemoryTaskBrokerOptions{
|
|
LeaseTTL: time.Second,
|
|
Now: func() time.Time { return now },
|
|
Material: TaskMaterialResolverFunc(func(context.Context, Candidate) (TaskMaterial, error) {
|
|
return TaskMaterial{Protocol: proxyDomain.SchemeHTTP, Host: "proxy.example", Port: 8080}, nil
|
|
}),
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewMemoryTaskBroker(): %v", err)
|
|
}
|
|
plan := PlannedTask{Candidate: Candidate{ProxyID: "proxy-a", State: proxyDomain.StateFetched, Level: healthDomain.LevelBasic, DueAt: now}, Deadline: now.Add(time.Minute), Attempts: 1}
|
|
if _, err := broker.Offer(context.Background(), []PlannedTask{plan}); err != nil {
|
|
t.Fatalf("Offer(): %v", err)
|
|
}
|
|
first, err := broker.Claim(context.Background(), TaskClaim{CheckerID: "checker-a", InstanceID: "instance-a", MaxInFlight: 1, SupportedLevels: []healthDomain.Level{healthDomain.LevelBasic}})
|
|
if err != nil || len(first) != 1 {
|
|
t.Fatalf("first Claim() = (%+v, %v)", first, err)
|
|
}
|
|
second, err := broker.Claim(context.Background(), TaskClaim{CheckerID: "checker-a", InstanceID: "instance-a", MaxInFlight: 1, SupportedLevels: []healthDomain.Level{healthDomain.LevelBasic}})
|
|
if err != nil || len(second) != 0 {
|
|
t.Fatalf("second Claim() = (%+v, %v)", second, err)
|
|
}
|
|
now = now.Add(2 * time.Second)
|
|
reclaimed, err := broker.Claim(context.Background(), TaskClaim{CheckerID: "checker-b", InstanceID: "instance-b", MaxInFlight: 1, SupportedLevels: []healthDomain.Level{healthDomain.LevelBasic}})
|
|
if err != nil || len(reclaimed) != 1 || reclaimed[0].TaskID != first[0].TaskID || reclaimed[0].LeaseToken == first[0].LeaseToken {
|
|
t.Fatalf("reclaimed Claim() = (%+v, %v)", reclaimed, err)
|
|
}
|
|
}
|
|
|
|
func TestMemoryTaskBrokerRejectsTerminalProxyState(t *testing.T) {
|
|
now := time.Date(2026, 8, 1, 11, 0, 0, 0, time.UTC)
|
|
broker, err := NewMemoryTaskBroker(MemoryTaskBrokerOptions{
|
|
LeaseTTL: time.Second, Now: func() time.Time { return now },
|
|
Material: TaskMaterialResolverFunc(func(context.Context, Candidate) (TaskMaterial, error) {
|
|
return TaskMaterial{Protocol: proxyDomain.SchemeHTTP, Host: "proxy.example", Port: 8080}, nil
|
|
}),
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewMemoryTaskBroker(): %v", err)
|
|
}
|
|
_, err = broker.Offer(context.Background(), []PlannedTask{{
|
|
Candidate: Candidate{ProxyID: "proxy-a", State: proxyDomain.StateExtracted, Level: healthDomain.LevelBasic, DueAt: now},
|
|
Deadline: now.Add(time.Second), Attempts: 1,
|
|
}})
|
|
if !errors.Is(err, ErrInvalidLeasedTask) {
|
|
t.Fatalf("Offer(terminal state) error = %v, want ErrInvalidLeasedTask", err)
|
|
}
|
|
}
|