nebula/timeout_test.go

139 lines
3.5 KiB
Go
Raw Normal View History

2019-11-19 10:00:20 -07:00
package nebula
import (
"github.com/stretchr/testify/assert"
"testing"
"time"
)
func TestNewTimerWheel(t *testing.T) {
// Make sure we get an object we expect
tw := NewTimerWheel(time.Second, time.Second*10)
assert.Equal(t, 11, tw.wheelLen)
assert.Equal(t, 0, tw.current)
assert.Nil(t, tw.lastTick)
assert.Equal(t, time.Second*1, tw.tickDuration)
assert.Equal(t, time.Second*10, tw.wheelDuration)
assert.Len(t, tw.wheel, 11)
// Assert the math is correct
tw = NewTimerWheel(time.Second*3, time.Second*10)
assert.Equal(t, 4, tw.wheelLen)
tw = NewTimerWheel(time.Second*120, time.Minute*10)
assert.Equal(t, 6, tw.wheelLen)
}
func TestTimerWheel_findWheel(t *testing.T) {
tw := NewTimerWheel(time.Second, time.Second*10)
assert.Len(t, tw.wheel, 11)
// Current + tick + 1 since we don't know how far into current we are
assert.Equal(t, 2, tw.findWheel(time.Second*1))
// Scale up to min duration
assert.Equal(t, 2, tw.findWheel(time.Millisecond*1))
// Make sure we hit that last index
assert.Equal(t, 0, tw.findWheel(time.Second*10))
// Scale down to max duration
assert.Equal(t, 0, tw.findWheel(time.Second*11))
tw.current = 1
// Make sure we account for the current position properly
assert.Equal(t, 3, tw.findWheel(time.Second*1))
assert.Equal(t, 1, tw.findWheel(time.Second*10))
}
func TestTimerWheel_Add(t *testing.T) {
tw := NewTimerWheel(time.Second, time.Second*10)
fp1 := FirewallPacket{}
tw.Add(fp1, time.Second*1)
// Make sure we set head and tail properly
assert.NotNil(t, tw.wheel[2])
assert.Equal(t, fp1, tw.wheel[2].Head.Packet)
assert.Nil(t, tw.wheel[2].Head.Next)
assert.Equal(t, fp1, tw.wheel[2].Tail.Packet)
assert.Nil(t, tw.wheel[2].Tail.Next)
// Make sure we only modify head
fp2 := FirewallPacket{}
tw.Add(fp2, time.Second*1)
assert.Equal(t, fp2, tw.wheel[2].Head.Packet)
assert.Equal(t, fp1, tw.wheel[2].Head.Next.Packet)
assert.Equal(t, fp1, tw.wheel[2].Tail.Packet)
assert.Nil(t, tw.wheel[2].Tail.Next)
// Make sure we use free'd items first
tw.itemCache = &TimeoutItem{}
tw.itemsCached = 1
tw.Add(fp2, time.Second*1)
assert.Nil(t, tw.itemCache)
assert.Equal(t, 0, tw.itemsCached)
}
func TestTimerWheel_Purge(t *testing.T) {
// First advance should set the lastTick and do nothing else
tw := NewTimerWheel(time.Second, time.Second*10)
assert.Nil(t, tw.lastTick)
tw.advance(time.Now())
assert.NotNil(t, tw.lastTick)
assert.Equal(t, 0, tw.current)
fps := []FirewallPacket{
{LocalIP: 1},
{LocalIP: 2},
{LocalIP: 3},
{LocalIP: 4},
}
tw.Add(fps[0], time.Second*1)
tw.Add(fps[1], time.Second*1)
tw.Add(fps[2], time.Second*2)
tw.Add(fps[3], time.Second*2)
ta := time.Now().Add(time.Second * 3)
lastTick := *tw.lastTick
tw.advance(ta)
assert.Equal(t, 3, tw.current)
assert.True(t, tw.lastTick.After(lastTick))
// Make sure we get all 4 packets back
for i := 0; i < 4; i++ {
p, has := tw.Purge()
assert.True(t, has)
assert.Equal(t, fps[i], p)
}
// Make sure there aren't any leftover
_, ok := tw.Purge()
assert.False(t, ok)
assert.Nil(t, tw.expired.Head)
assert.Nil(t, tw.expired.Tail)
// Make sure we cached the free'd items
assert.Equal(t, 4, tw.itemsCached)
ci := tw.itemCache
for i := 0; i < 4; i++ {
assert.NotNil(t, ci)
ci = ci.Next
}
assert.Nil(t, ci)
// Lets make sure we roll over properly
ta = ta.Add(time.Second * 5)
tw.advance(ta)
assert.Equal(t, 8, tw.current)
ta = ta.Add(time.Second * 2)
tw.advance(ta)
assert.Equal(t, 10, tw.current)
ta = ta.Add(time.Second * 1)
tw.advance(ta)
assert.Equal(t, 0, tw.current)
}