tailscale/net/flowtrack/flowtrack_test.go

129 lines
2.8 KiB
Go

// Copyright (c) Tailscale Inc & AUTHORS
// SPDX-License-Identifier: BSD-3-Clause
package flowtrack
import (
"encoding/json"
"net/netip"
"testing"
"tailscale.com/tstest"
"tailscale.com/types/ipproto"
)
func TestCache(t *testing.T) {
c := &Cache[int]{MaxEntries: 2}
k1 := MakeTuple(0, netip.MustParseAddrPort("1.1.1.1:1"), netip.MustParseAddrPort("1.1.1.1:1"))
k2 := MakeTuple(0, netip.MustParseAddrPort("1.1.1.1:1"), netip.MustParseAddrPort("2.2.2.2:2"))
k3 := MakeTuple(0, netip.MustParseAddrPort("1.1.1.1:1"), netip.MustParseAddrPort("3.3.3.3:3"))
k4 := MakeTuple(0, netip.MustParseAddrPort("1.1.1.1:1"), netip.MustParseAddrPort("4.4.4.4:4"))
wantLen := func(want int) {
t.Helper()
if got := c.Len(); got != want {
t.Fatalf("Len = %d; want %d", got, want)
}
}
wantVal := func(key Tuple, want int) {
t.Helper()
got, ok := c.Get(key)
if !ok {
t.Fatalf("Get(%q) failed; want value %v", key, want)
}
if *got != want {
t.Fatalf("Get(%q) = %v; want %v", key, got, want)
}
}
wantMissing := func(key Tuple) {
t.Helper()
if got, ok := c.Get(key); ok {
t.Fatalf("Get(%q) = %v; want absent from cache", key, got)
}
}
wantLen(0)
c.RemoveOldest() // shouldn't panic
c.Remove(k4) // shouldn't panic
c.Add(k1, 1)
wantLen(1)
c.Add(k2, 2)
wantLen(2)
c.Add(k3, 3)
wantLen(2) // hit the max
wantMissing(k1)
c.Remove(k1)
wantLen(2) // no change; k1 should've been the deleted one per LRU
wantVal(k3, 3)
wantVal(k2, 2)
c.Remove(k2)
wantLen(1)
wantMissing(k2)
c.Add(k3, 30)
wantVal(k3, 30)
wantLen(1)
err := tstest.MinAllocsPerRun(t, 0, func() {
got, ok := c.Get(k3)
if !ok {
t.Fatal("missing k3")
}
if *got != 30 {
t.Fatalf("got = %d; want 30", got)
}
})
if err != nil {
t.Error(err)
}
}
func BenchmarkMapKeys(b *testing.B) {
b.Run("typed", func(b *testing.B) {
c := &Cache[struct{}]{MaxEntries: 1000}
var t Tuple
for proto := range 20 {
t = Tuple{proto: ipproto.Proto(proto), src: netip.MustParseAddr("1.1.1.1").As16(), srcPort: 1, dst: netip.MustParseAddr("1.1.1.1").As16(), dstPort: 1}
c.Add(t, struct{}{})
}
for i := 0; i < b.N; i++ {
_, ok := c.Get(t)
if !ok {
b.Fatal("missing key")
}
}
})
}
func TestStringJSON(t *testing.T) {
v := MakeTuple(123,
netip.MustParseAddrPort("1.2.3.4:5"),
netip.MustParseAddrPort("6.7.8.9:10"))
if got, want := v.String(), "(IPProto-123 1.2.3.4:5 => 6.7.8.9:10)"; got != want {
t.Errorf("String = %q; want %q", got, want)
}
got, err := json.Marshal(v)
if err != nil {
t.Fatal(err)
}
const want = `{"proto":123,"src":"1.2.3.4:5","dst":"6.7.8.9:10"}`
if string(got) != want {
t.Errorf("Marshal = %q; want %q", got, want)
}
var back Tuple
if err := json.Unmarshal(got, &back); err != nil {
t.Fatal(err)
}
if back != v {
t.Errorf("back = %v; want %v", back, v)
}
}