123 lines
2.2 KiB
Go
123 lines
2.2 KiB
Go
package pubsub
|
|
|
|
import (
|
|
"context"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/libp2p/go-libp2p-core/peer"
|
|
)
|
|
|
|
func TestPeerGater(t *testing.T) {
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
defer cancel()
|
|
|
|
peerA := peer.ID("A")
|
|
peerAip := "1.2.3.4"
|
|
|
|
pg := newPeerGater(ctx, nil, .1, .9)
|
|
pg.getIP = func(p peer.ID) string {
|
|
switch p {
|
|
case peerA:
|
|
return peerAip
|
|
default:
|
|
return "<wtf>"
|
|
}
|
|
}
|
|
|
|
pg.AddPeer(peerA, "")
|
|
|
|
status := pg.AcceptFrom(peerA)
|
|
if status != AcceptAll {
|
|
t.Fatal("expected AcceptAll")
|
|
}
|
|
|
|
msg := &Message{ReceivedFrom: peerA}
|
|
|
|
pg.ValidateMessage(msg)
|
|
status = pg.AcceptFrom(peerA)
|
|
if status != AcceptAll {
|
|
t.Fatal("expected AcceptAll")
|
|
}
|
|
|
|
pg.RejectMessage(msg, rejectValidationQueueFull)
|
|
status = pg.AcceptFrom(peerA)
|
|
if status != AcceptAll {
|
|
t.Fatal("expected AcceptAll")
|
|
}
|
|
|
|
pg.RejectMessage(msg, rejectValidationThrottled)
|
|
status = pg.AcceptFrom(peerA)
|
|
if status != AcceptAll {
|
|
t.Fatal("expected AcceptAll")
|
|
}
|
|
|
|
for i := 0; i < 100; i++ {
|
|
pg.RejectMessage(msg, rejectValidationIgnored)
|
|
pg.RejectMessage(msg, rejectValidationFailed)
|
|
}
|
|
|
|
accepted := false
|
|
for i := 0; !accepted && i < 1000; i++ {
|
|
status = pg.AcceptFrom(peerA)
|
|
if status == AcceptControl {
|
|
accepted = true
|
|
}
|
|
}
|
|
if !accepted {
|
|
t.Fatal("expected AcceptControl")
|
|
}
|
|
|
|
for i := 0; i < 100; i++ {
|
|
pg.DeliverMessage(msg)
|
|
}
|
|
|
|
accepted = false
|
|
for i := 0; !accepted && i < 1000; i++ {
|
|
status = pg.AcceptFrom(peerA)
|
|
if status == AcceptAll {
|
|
accepted = true
|
|
}
|
|
}
|
|
if !accepted {
|
|
t.Fatal("expected to accept at least once")
|
|
}
|
|
|
|
for i := 0; i < 100; i++ {
|
|
pg.decayStats()
|
|
}
|
|
|
|
status = pg.AcceptFrom(peerA)
|
|
if status != AcceptAll {
|
|
t.Fatal("expected AcceptAll")
|
|
}
|
|
|
|
pg.RemovePeer(peerA)
|
|
pg.Lock()
|
|
_, ok := pg.peerStats[peerA]
|
|
pg.Unlock()
|
|
if ok {
|
|
t.Fatal("still have a stat record for peerA")
|
|
}
|
|
|
|
pg.Lock()
|
|
_, ok = pg.ipStats[peerAip]
|
|
pg.Unlock()
|
|
if !ok {
|
|
t.Fatal("expected to still have a stat record for peerA's ip")
|
|
}
|
|
|
|
pg.Lock()
|
|
pg.ipStats[peerAip].expire = time.Now()
|
|
pg.Unlock()
|
|
|
|
time.Sleep(2 * time.Second)
|
|
|
|
pg.Lock()
|
|
_, ok = pg.ipStats["1.2.3.4"]
|
|
pg.Unlock()
|
|
if ok {
|
|
t.Fatal("still have a stat record for peerA's ip")
|
|
}
|
|
}
|