247 lines
7.4 KiB
Go
247 lines
7.4 KiB
Go
package extraction
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
func TestMemoryStoreNeverExtractsProxyTwice(t *testing.T) {
|
|
now := time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC)
|
|
store := NewMemoryStore([]Candidate{
|
|
{ID: "p1", State: Available, ExpiresAt: now.Add(time.Minute), LastCheckedAt: now},
|
|
})
|
|
|
|
var wg sync.WaitGroup
|
|
results := make(chan string, 1000)
|
|
for range 1000 {
|
|
wg.Add(1)
|
|
go func() {
|
|
defer wg.Done()
|
|
result, err := store.Extract(context.Background(), Command{
|
|
Requested: 1,
|
|
Fulfillment: Partial,
|
|
Now: now,
|
|
MinRemainingTTL: 30 * time.Second,
|
|
MaxHealthCheckAge: 10 * time.Second,
|
|
})
|
|
if err != nil {
|
|
t.Errorf("Extract(): %v", err)
|
|
return
|
|
}
|
|
for _, item := range result.Items {
|
|
results <- item.ID
|
|
}
|
|
}()
|
|
}
|
|
wg.Wait()
|
|
close(results)
|
|
|
|
count := 0
|
|
for id := range results {
|
|
if id != "p1" {
|
|
t.Fatalf("unexpected proxy %q", id)
|
|
}
|
|
count++
|
|
}
|
|
if count != 1 {
|
|
t.Fatalf("proxy extracted %d times, want exactly once", count)
|
|
}
|
|
}
|
|
|
|
func TestAllOrNothingDoesNotConsumePartialInventory(t *testing.T) {
|
|
now := time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC)
|
|
store := NewMemoryStore([]Candidate{
|
|
{ID: "p1", State: Available, ExpiresAt: now.Add(time.Minute), LastCheckedAt: now},
|
|
})
|
|
|
|
result, err := store.Extract(context.Background(), Command{
|
|
Requested: 2,
|
|
Fulfillment: AllOrNothing,
|
|
Now: now,
|
|
MinRemainingTTL: 30 * time.Second,
|
|
MaxHealthCheckAge: 10 * time.Second,
|
|
})
|
|
if err != ErrInsufficientProxies {
|
|
t.Fatalf("Extract() error = %v, want ErrInsufficientProxies", err)
|
|
}
|
|
if len(result.Items) != 0 {
|
|
t.Fatalf("Extract() returned %d items, want 0", len(result.Items))
|
|
}
|
|
|
|
partial, err := store.Extract(context.Background(), Command{
|
|
Requested: 1,
|
|
Fulfillment: Partial,
|
|
Now: now,
|
|
MinRemainingTTL: 30 * time.Second,
|
|
MaxHealthCheckAge: 10 * time.Second,
|
|
})
|
|
if err != nil || len(partial.Items) != 1 {
|
|
t.Fatalf("inventory was consumed by failed all-or-nothing: result=%+v err=%v", partial, err)
|
|
}
|
|
}
|
|
|
|
func TestMemoryStoreCommitsAuditWithExtraction(t *testing.T) {
|
|
now := time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC)
|
|
store := NewMemoryStore([]Candidate{{
|
|
ID: "p1",
|
|
Upstream: "provider-a",
|
|
State: Available,
|
|
ExpiresAt: now.Add(time.Minute),
|
|
LastCheckedAt: now,
|
|
}})
|
|
|
|
_, err := store.Extract(context.Background(), Command{
|
|
RequestID: "req-1",
|
|
ClientID: "client-1",
|
|
SourceIP: "192.0.2.30",
|
|
Requested: 1,
|
|
Fulfillment: Partial,
|
|
Now: now,
|
|
MinRemainingTTL: 30 * time.Second,
|
|
MaxHealthCheckAge: 10 * time.Second,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Extract(): %v", err)
|
|
}
|
|
|
|
records := store.Records()
|
|
if len(records) != 1 {
|
|
t.Fatalf("audit records = %d, want 1", len(records))
|
|
}
|
|
record := records[0]
|
|
if record.ProxyID != "p1" || record.ClientID != "client-1" || record.RequestID != "req-1" {
|
|
t.Fatalf("audit record = %+v", record)
|
|
}
|
|
if !record.ExtractedAt.Equal(now) || !record.ExpiresAt.Equal(now.Add(time.Minute)) {
|
|
t.Fatalf("audit timestamps = %+v", record)
|
|
}
|
|
}
|
|
|
|
func TestMemoryStoreReplaysCommittedIdempotentResult(t *testing.T) {
|
|
now := time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC)
|
|
store := NewMemoryStore([]Candidate{{
|
|
ID: "p1",
|
|
State: Available,
|
|
ExpiresAt: now.Add(time.Minute),
|
|
LastCheckedAt: now,
|
|
}})
|
|
command := Command{
|
|
RequestID: "req-1",
|
|
ClientID: "client-1",
|
|
IdempotencyKey: "idem-12345678",
|
|
Requested: 1,
|
|
Fulfillment: Partial,
|
|
Now: now,
|
|
MinRemainingTTL: 30 * time.Second,
|
|
MaxHealthCheckAge: 10 * time.Second,
|
|
}
|
|
|
|
first, err := store.Extract(context.Background(), command)
|
|
if err != nil {
|
|
t.Fatalf("first Extract(): %v", err)
|
|
}
|
|
command.RequestID = "req-2"
|
|
second, err := store.Extract(context.Background(), command)
|
|
if err != nil {
|
|
t.Fatalf("second Extract(): %v", err)
|
|
}
|
|
|
|
if len(first.Items) != 1 || len(second.Items) != 1 || second.Items[0].ID != first.Items[0].ID {
|
|
t.Fatalf("idempotent results: first=%+v second=%+v", first, second)
|
|
}
|
|
if got := len(store.Records()); got != 1 {
|
|
t.Fatalf("audit records = %d, want 1", got)
|
|
}
|
|
}
|
|
|
|
func TestMemoryStoreRejectsIdempotencyKeyReuseWithDifferentRequest(t *testing.T) {
|
|
now := time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC)
|
|
store := NewMemoryStore([]Candidate{{
|
|
ID: "p1",
|
|
State: Available,
|
|
ExpiresAt: now.Add(time.Minute),
|
|
LastCheckedAt: now,
|
|
}})
|
|
command := Command{
|
|
RequestID: "req-1",
|
|
ClientID: "client-1",
|
|
IdempotencyKey: "idem-12345678",
|
|
Requested: 1,
|
|
Fulfillment: Partial,
|
|
Now: now,
|
|
MinRemainingTTL: 30 * time.Second,
|
|
MaxHealthCheckAge: 10 * time.Second,
|
|
}
|
|
if _, err := store.Extract(context.Background(), command); err != nil {
|
|
t.Fatalf("first Extract(): %v", err)
|
|
}
|
|
|
|
command.Requested = 2
|
|
if _, err := store.Extract(context.Background(), command); err != ErrIdempotencyConflict {
|
|
t.Fatalf("second Extract() error = %v, want ErrIdempotencyConflict", err)
|
|
}
|
|
}
|
|
|
|
func TestMemoryStoreExtractsOnlyUnownedProxy(t *testing.T) {
|
|
now := time.Date(2026, 7, 28, 10, 0, 0, 0, time.UTC)
|
|
store := NewMemoryStore([]Candidate{
|
|
{ID: "owned", OwnerWorkerID: "worker-1", State: Available, ExpiresAt: now.Add(time.Minute), LastCheckedAt: now},
|
|
{ID: "unowned", State: Available, ExpiresAt: now.Add(time.Minute), LastCheckedAt: now},
|
|
})
|
|
|
|
result, err := store.Extract(context.Background(), Command{
|
|
RequestID: "req-1",
|
|
ClientID: "client-1",
|
|
Requested: 2,
|
|
Fulfillment: Partial,
|
|
Now: now,
|
|
MinRemainingTTL: 30 * time.Second,
|
|
MaxHealthCheckAge: 10 * time.Second,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Extract(): %v", err)
|
|
}
|
|
if len(result.Items) != 1 || result.Items[0].ID != "unowned" {
|
|
t.Fatalf("extracted items = %+v, want only unowned", result.Items)
|
|
}
|
|
}
|
|
|
|
func TestMemoryStoreValidatesCommandAndHonorsCancellation(t *testing.T) {
|
|
store := NewMemoryStore([]Candidate{{ID: "p1", State: Available}})
|
|
if _, err := store.Extract(context.Background(), Command{
|
|
Requested: 1, Fulfillment: Partial, ReserveForGateway: -1,
|
|
}); !errors.Is(err, ErrInvalidCommand) {
|
|
t.Fatalf("Extract(negative reserve) error = %v, want ErrInvalidCommand", err)
|
|
}
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
cancel()
|
|
if _, err := store.Extract(ctx, Command{Requested: 1, Fulfillment: Partial}); !errors.Is(err, context.Canceled) {
|
|
t.Fatalf("Extract(canceled) error = %v, want context.Canceled", err)
|
|
}
|
|
result, err := store.Extract(context.Background(), Command{Requested: 1, Fulfillment: Partial})
|
|
if err != nil || result.Returned != 1 {
|
|
t.Fatalf("candidate changed after canceled request: result=%+v err=%v", result, err)
|
|
}
|
|
}
|
|
|
|
func TestMemoryStoreClonesIdempotencyCommandFilters(t *testing.T) {
|
|
store := NewMemoryStore([]Candidate{{ID: "p1", Protocol: "http", State: Available}})
|
|
protocols := []string{"http"}
|
|
command := Command{
|
|
ClientID: "client-1", IdempotencyKey: "idem-1", Requested: 1,
|
|
Fulfillment: Partial, Protocols: protocols,
|
|
}
|
|
if _, err := store.Extract(context.Background(), command); err != nil {
|
|
t.Fatalf("first Extract(): %v", err)
|
|
}
|
|
protocols[0] = "socks5"
|
|
command.Protocols = []string{"http"}
|
|
if _, err := store.Extract(context.Background(), command); err != nil {
|
|
t.Fatalf("idempotent replay after caller mutation: %v", err)
|
|
}
|
|
}
|