250 lines
7.1 KiB
Go
250 lines
7.1 KiB
Go
package extraction
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
|
|
"proxy-pool/internal/domain/activitypool"
|
|
domain "proxy-pool/internal/domain/extraction"
|
|
proxyDomain "proxy-pool/internal/domain/proxy"
|
|
)
|
|
|
|
func TestServiceAppliesPolicyAndBuildsResponse(t *testing.T) {
|
|
now := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
|
|
store := &recordingStore{result: domain.Result{
|
|
Requested: 2,
|
|
Returned: 1,
|
|
Items: []domain.Candidate{{
|
|
ID: "p1",
|
|
Protocol: "http",
|
|
Host: "192.0.2.10",
|
|
Port: 8080,
|
|
URL: "http://user:pass@192.0.2.10:8080",
|
|
Upstream: "provider-a",
|
|
ExpiresAt: now.Add(95 * time.Second),
|
|
}},
|
|
}}
|
|
service, err := NewService(store, Policy{
|
|
MaxCountPerRequest: 20,
|
|
DefaultFulfillment: domain.Partial,
|
|
MinRemainingTTL: 30 * time.Second,
|
|
MaxHealthCheckAge: 15 * time.Second,
|
|
ReserveForGateway: 5,
|
|
IdempotencyTTL: 2 * time.Minute,
|
|
}, allowAllAdmission{}, func() time.Time { return now })
|
|
if err != nil {
|
|
t.Fatalf("NewService(): %v", err)
|
|
}
|
|
|
|
response, err := service.Extract(context.Background(), Request{
|
|
RequestID: "req-1",
|
|
ClientID: "client-1",
|
|
SourceIP: "192.0.2.30",
|
|
Count: 2,
|
|
Filters: Filters{
|
|
Protocols: []string{"http"},
|
|
Regions: []string{"shanghai"},
|
|
},
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Extract(): %v", err)
|
|
}
|
|
|
|
if store.command.Fulfillment != domain.Partial || store.command.ReserveForGateway != 5 {
|
|
t.Fatalf("store command policy = %+v", store.command)
|
|
}
|
|
if store.command.RequestID != "req-1" || store.command.ClientID != "client-1" {
|
|
t.Fatalf("store command request context = %+v", store.command)
|
|
}
|
|
if store.command.IdempotencyTTL != 2*time.Minute {
|
|
t.Fatalf("store command idempotency TTL = %s, want 2m", store.command.IdempotencyTTL)
|
|
}
|
|
if response.RequestID != "req-1" || response.Requested != 2 || response.Returned != 1 {
|
|
t.Fatalf("response = %+v", response)
|
|
}
|
|
if len(response.Proxies) != 1 || response.Proxies[0].RemainingTTLSeconds != 95 {
|
|
t.Fatalf("response proxies = %+v", response.Proxies)
|
|
}
|
|
if !response.Proxies[0].ExtractedAt.Equal(now) {
|
|
t.Fatalf("extractedAt = %s, want %s", response.Proxies[0].ExtractedAt, now)
|
|
}
|
|
}
|
|
|
|
func TestServiceRejectsCountAbovePolicyBeforeStore(t *testing.T) {
|
|
store := &recordingStore{}
|
|
service, err := NewService(store, Policy{
|
|
MaxCountPerRequest: 1,
|
|
DefaultFulfillment: domain.Partial,
|
|
}, allowAllAdmission{}, time.Now)
|
|
if err != nil {
|
|
t.Fatalf("NewService(): %v", err)
|
|
}
|
|
|
|
_, err = service.Extract(context.Background(), Request{
|
|
RequestID: "req-1",
|
|
ClientID: "client-1",
|
|
Count: 2,
|
|
})
|
|
if !errors.Is(err, ErrCountExceeded) {
|
|
t.Fatalf("Extract() error = %v, want ErrCountExceeded", err)
|
|
}
|
|
if store.calls != 0 {
|
|
t.Fatalf("store calls = %d, want 0", store.calls)
|
|
}
|
|
}
|
|
|
|
func TestServiceAppliesAdmissionBeforeStoreUsingStableIdentity(t *testing.T) {
|
|
store := &recordingStore{}
|
|
admission := &recordingAdmission{err: ErrAdmissionRejected}
|
|
service, err := NewService(store, Policy{
|
|
MaxCountPerRequest: 1,
|
|
DefaultFulfillment: domain.Partial,
|
|
}, admission, time.Now)
|
|
if err != nil {
|
|
t.Fatalf("NewService(): %v", err)
|
|
}
|
|
|
|
_, err = service.Extract(context.Background(), Request{
|
|
RequestID: "req-1",
|
|
SourceIP: "192.0.2.30",
|
|
Count: 1,
|
|
})
|
|
if !errors.Is(err, ErrAdmissionRejected) {
|
|
t.Fatalf("Extract() error = %v, want ErrAdmissionRejected", err)
|
|
}
|
|
if admission.key != "source:192.0.2.30" || admission.calls != 1 {
|
|
t.Fatalf("admission = %+v, want source identity", admission)
|
|
}
|
|
if store.calls != 0 {
|
|
t.Fatalf("store calls = %d, want 0", store.calls)
|
|
}
|
|
}
|
|
|
|
func TestServiceUsesSourceIdentityForEphemeralIdempotency(t *testing.T) {
|
|
store := &recordingStore{}
|
|
service, err := NewService(store, Policy{
|
|
MaxCountPerRequest: 1,
|
|
DefaultFulfillment: domain.Partial,
|
|
}, allowAllAdmission{}, time.Now)
|
|
if err != nil {
|
|
t.Fatalf("NewService(): %v", err)
|
|
}
|
|
|
|
_, err = service.Extract(context.Background(), Request{
|
|
RequestID: "req-1",
|
|
SourceIP: "192.0.2.30",
|
|
Count: 1,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Extract(): %v", err)
|
|
}
|
|
if store.command.ClientID != "192.0.2.30" {
|
|
t.Fatalf("store command ClientID = %q, want source IP", store.command.ClientID)
|
|
}
|
|
}
|
|
|
|
func TestServiceClassifiesStoreUnavailableAndPreservesCause(t *testing.T) {
|
|
t.Parallel()
|
|
storeErr := errors.Join(domain.ErrStoreUnavailable, errors.New("redis dial failed: TOKEN"))
|
|
service, err := NewService(&recordingStore{err: storeErr}, Policy{
|
|
MaxCountPerRequest: 1,
|
|
DefaultFulfillment: domain.Partial,
|
|
}, allowAllAdmission{}, time.Now)
|
|
if err != nil {
|
|
t.Fatalf("NewService(): %v", err)
|
|
}
|
|
|
|
_, err = service.Extract(context.Background(), Request{
|
|
RequestID: "req-1",
|
|
ClientID: "client-1",
|
|
Count: 1,
|
|
})
|
|
if !errors.Is(err, ErrUnavailable) || !errors.Is(err, domain.ErrStoreUnavailable) {
|
|
t.Fatalf("Extract() error = %v, want unavailable classification with store cause", err)
|
|
}
|
|
}
|
|
|
|
func TestServiceIdempotentReplayKeepsOriginalExtractionTime(t *testing.T) {
|
|
firstTime := time.Date(2026, 7, 28, 12, 0, 0, 0, time.UTC)
|
|
secondTime := firstTime.Add(time.Minute)
|
|
store := activitypool.NewMemoryPool()
|
|
checkedAt := firstTime
|
|
upserted, err := store.UpsertFetched(context.Background(), "provider-a", activitypool.FetchedBatch{
|
|
ObservedAt: firstTime,
|
|
ConfiguredTTL: 2 * time.Minute,
|
|
MaxSize: 100,
|
|
Proxies: []proxyDomain.Proxy{{
|
|
ID: "p1", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080,
|
|
State: proxyDomain.StateAvailable, LastCheckedAt: &checkedAt,
|
|
}},
|
|
})
|
|
if err != nil || upserted.Inserted != 1 {
|
|
t.Fatalf("UpsertFetched() = %+v, %v", upserted, err)
|
|
}
|
|
times := []time.Time{firstTime, secondTime}
|
|
service, err := NewService(store, Policy{
|
|
MaxCountPerRequest: 1,
|
|
DefaultFulfillment: domain.Partial,
|
|
}, allowAllAdmission{}, func() time.Time {
|
|
value := times[0]
|
|
times = times[1:]
|
|
return value
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("NewService(): %v", err)
|
|
}
|
|
request := Request{
|
|
RequestID: "req-1",
|
|
ClientID: "client-1",
|
|
IdempotencyKey: "idem-12345678",
|
|
Count: 1,
|
|
}
|
|
|
|
first, err := service.Extract(context.Background(), request)
|
|
if err != nil {
|
|
t.Fatalf("first Extract(): %v", err)
|
|
}
|
|
request.RequestID = "req-2"
|
|
second, err := service.Extract(context.Background(), request)
|
|
if err != nil {
|
|
t.Fatalf("second Extract(): %v", err)
|
|
}
|
|
if !second.Proxies[0].ExtractedAt.Equal(first.Proxies[0].ExtractedAt) {
|
|
t.Fatalf("replayed extractedAt = %s, want %s", second.Proxies[0].ExtractedAt, first.Proxies[0].ExtractedAt)
|
|
}
|
|
if second.Proxies[0].RemainingTTLSeconds != first.Proxies[0].RemainingTTLSeconds {
|
|
t.Fatalf("replayed remaining TTL = %d, want %d", second.Proxies[0].RemainingTTLSeconds, first.Proxies[0].RemainingTTLSeconds)
|
|
}
|
|
}
|
|
|
|
type recordingStore struct {
|
|
command domain.Command
|
|
result domain.Result
|
|
err error
|
|
calls int
|
|
}
|
|
|
|
type recordingAdmission struct {
|
|
key string
|
|
err error
|
|
calls int
|
|
}
|
|
|
|
type allowAllAdmission struct{}
|
|
|
|
func (allowAllAdmission) Admit(context.Context, string) error { return nil }
|
|
|
|
func (a *recordingAdmission) Admit(_ context.Context, key string) error {
|
|
a.calls++
|
|
a.key = key
|
|
return a.err
|
|
}
|
|
|
|
func (s *recordingStore) Extract(_ context.Context, command domain.Command) (domain.Result, error) {
|
|
s.calls++
|
|
s.command = command
|
|
return s.result, s.err
|
|
}
|