108 lines
4.2 KiB
Go
108 lines
4.2 KiB
Go
package worker
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"testing"
|
|
"time"
|
|
|
|
"proxy-pool/internal/controlplane/snapshotwire"
|
|
ownershipDomain "proxy-pool/internal/domain/ownership"
|
|
proxyDomain "proxy-pool/internal/domain/proxy"
|
|
)
|
|
|
|
func TestInitialSnapshotSourceIssuesNextFullSnapshot(t *testing.T) {
|
|
now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
|
|
source, err := NewInitialSnapshotSource(epochReaderStub{epoch: 9}, time.Minute, func() time.Time { return now })
|
|
if err != nil {
|
|
t.Fatalf("NewInitialSnapshotSource(): %v", err)
|
|
}
|
|
updates, err := source.Watch(context.Background(), SnapshotWatchRequest{
|
|
WorkerID: "worker-a", SessionID: "session-a", LastAppliedVersion: 4,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("Watch(): %v", err)
|
|
}
|
|
full := <-updates
|
|
if full.GetVersion() != 5 || full.GetOwnershipEpoch() != 9 || !full.GetValidUntil().AsTime().Equal(now.Add(time.Minute)) {
|
|
t.Fatalf("snapshot = %+v", full)
|
|
}
|
|
checksum, err := snapshotwire.Checksum(full)
|
|
if err != nil || string(checksum[:]) != string(full.GetChecksum()) {
|
|
t.Fatalf("snapshot checksum = %x, %v; want %x", full.GetChecksum(), err, checksum)
|
|
}
|
|
}
|
|
|
|
func TestOwnedSnapshotSourceBuildsBoundedProxySnapshot(t *testing.T) {
|
|
now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
|
|
expiresAt := now.Add(10 * time.Minute)
|
|
usableUntil := now.Add(5 * time.Minute)
|
|
source, err := NewOwnedSnapshotSource(epochReaderStub{epoch: 9}, snapshotReaderStub{proxies: []ownershipDomain.SnapshotProxy{{
|
|
Proxy: proxyDomain.Proxy{
|
|
ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080,
|
|
SourceUpstream: "upstream-a", ExpiresAt: &expiresAt, UsableUntil: &usableUntil,
|
|
MaxConcurrency: 7, State: proxyDomain.StateAvailable, Tags: map[string]string{"region": "cn"},
|
|
},
|
|
OwnershipEpoch: 4, LeaseExpiresAt: now.Add(time.Minute),
|
|
}}}, time.Minute*2, 10, 4096, func() time.Time { return now })
|
|
if err != nil {
|
|
t.Fatalf("NewOwnedSnapshotSource(): %v", err)
|
|
}
|
|
updates, err := source.Watch(context.Background(), SnapshotWatchRequest{WorkerID: "worker-a", SessionID: "session-a"})
|
|
if err != nil {
|
|
t.Fatalf("Watch(): %v", err)
|
|
}
|
|
full := <-updates
|
|
if full.GetVersion() != 1 || full.GetOwnershipEpoch() != 9 || len(full.GetProxies()) != 1 {
|
|
t.Fatalf("snapshot = %+v", full)
|
|
}
|
|
proxy := full.GetProxies()[0]
|
|
if proxy.GetOwnershipEpoch() != 4 || !proxy.GetUsableUntil().AsTime().Equal(now.Add(time.Minute)) ||
|
|
!full.GetValidUntil().AsTime().Equal(now.Add(time.Minute)) {
|
|
t.Fatalf("wire proxy = %+v, valid until = %s", proxy, full.GetValidUntil().AsTime())
|
|
}
|
|
checksum, err := snapshotwire.Checksum(full)
|
|
if err != nil || string(checksum[:]) != string(full.GetChecksum()) {
|
|
t.Fatalf("snapshot checksum = %x, %v; want %x", full.GetChecksum(), err, checksum)
|
|
}
|
|
}
|
|
|
|
func TestOwnedSnapshotSourceRejectsCredentialReferenceUntilMaterialIsAvailable(t *testing.T) {
|
|
now := time.Date(2026, 7, 31, 12, 0, 0, 0, time.UTC)
|
|
expiresAt := now.Add(time.Minute)
|
|
usableUntil := now.Add(time.Minute)
|
|
source, err := NewOwnedSnapshotSource(epochReaderStub{epoch: 1}, snapshotReaderStub{proxies: []ownershipDomain.SnapshotProxy{{
|
|
Proxy: proxyDomain.Proxy{
|
|
ID: "proxy-a", Scheme: proxyDomain.SchemeHTTP, Host: "192.0.2.10", Port: 8080,
|
|
SourceUpstream: "upstream-a", ExpiresAt: &expiresAt, UsableUntil: &usableUntil,
|
|
MaxConcurrency: 1, State: proxyDomain.StateAvailable, CredentialVersion: "v1",
|
|
},
|
|
OwnershipEpoch: 1, LeaseExpiresAt: now.Add(time.Minute),
|
|
}}}, time.Minute, 10, 4096, func() time.Time { return now })
|
|
if err != nil {
|
|
t.Fatalf("NewOwnedSnapshotSource(): %v", err)
|
|
}
|
|
_, err = source.Watch(context.Background(), SnapshotWatchRequest{WorkerID: "worker-a", SessionID: "session-a"})
|
|
if !errors.Is(err, ErrSnapshotCredentialsUnavailable) {
|
|
t.Fatalf("Watch() error = %v, want ErrSnapshotCredentialsUnavailable", err)
|
|
}
|
|
}
|
|
|
|
type epochReaderStub struct {
|
|
epoch uint64
|
|
err error
|
|
}
|
|
|
|
func (reader epochReaderStub) CurrentOwnershipEpoch(context.Context) (uint64, error) {
|
|
return reader.epoch, reader.err
|
|
}
|
|
|
|
type snapshotReaderStub struct {
|
|
proxies []ownershipDomain.SnapshotProxy
|
|
err error
|
|
}
|
|
|
|
func (reader snapshotReaderStub) ReadWorkerSnapshot(context.Context, string, int) ([]ownershipDomain.SnapshotProxy, error) {
|
|
return reader.proxies, reader.err
|
|
}
|