diff --git a/probe/endpoint/connection_tracker_test.go b/probe/endpoint/connection_tracker_test.go new file mode 100644 index 000000000..1b1c82f31 --- /dev/null +++ b/probe/endpoint/connection_tracker_test.go @@ -0,0 +1,112 @@ +package endpoint + +import ( + "testing" +) + +// mapPortToPids collects info about connections between specific +// address pairs and destination port. + +// a set of connections from a single pid to a destination off-box (pid is zero) +func fakeConnectionsOffBox(count int, fromPid uint32, startPort uint16) mapPortToPids { + ret := make(mapPortToPids, count) + for i := 0; i < count; i++ { + ret[uint16(i)+startPort] = pidPair{fromPid: fromPid, toPid: 0} + } + return ret +} + +// a set of connections to a single pid from a destination off-box (pid is zero) +func fakeConnectionsFromOffBox(count int, toPid uint32, startPort uint16) mapPortToPids { + ret := make(mapPortToPids, count) + for i := 0; i < count; i++ { + ret[uint16(i)+startPort] = pidPair{fromPid: 0, toPid: toPid} + } + return ret +} + +// a set of connections from a range of pids and a range of ports between two pids +func fakeConnections(count int, fromPid, toPid uint32, startPort uint16) mapPortToPids { + ret := make(mapPortToPids, count) + for i := 0; i < count; i++ { + ret[uint16(i)+startPort] = pidPair{fromPid: fromPid, toPid: toPid} + } + return ret +} + +// union N existing sets +func concatMapPortToPids(maps ...mapPortToPids) mapPortToPids { + ret := make(mapPortToPids) + for _, m := range maps { + for port, pair := range m { + ret[port] = pair + } + } + return ret +} + +func TestConnectionThinning(t *testing.T) { + for _, d := range []struct { + name string + m mapPortToPids + expC int + }{ + { + name: "0 ports", + m: mapPortToPids{}, + expC: 0, + }, + { + name: "5 connections off-box", + m: fakeConnectionsOffBox(5, 1000, 30000), + expC: 5, // 5 is too few to thin down + }, + { + name: "50 connections off-box", + m: fakeConnectionsOffBox(50, 1000, 30000), + expC: 5, // 50 connections should be thinned down + }, + { + name: "50 connections from off-box", + m: fakeConnectionsFromOffBox(50, 1000, 30000), + expC: 5, // 50 connections should be thinned down + }, + { + name: "5 connections from pid 1000 to 2000", + m: fakeConnections(5, 1000, 2000, 30000), + expC: 5, // 5 is too few to thin down + }, + { + name: "50 connections from pid 1000 to 2000", + m: fakeConnections(50, 1000, 2000, 30000), + expC: 5, // 50 connections should be thinned down + }, + { + name: "connections both ways", + m: concatMapPortToPids( + fakeConnections(50, 1000, 2000, 30000), + fakeConnections(50, 2000, 1000, 40000)), + expC: 100, // no thinning because in both directions + }, + { + name: "connections to two different pids", + m: concatMapPortToPids( + fakeConnections(50, 1000, 2000, 30000), + fakeConnections(50, 1000, 3000, 40000)), + expC: 100, // no thinning because two different pids + }, + { + name: "connections from off-box and on-box", + m: concatMapPortToPids(fakeConnectionsFromOffBox(50, 1000, 30000), fakeConnections(50, 2000, 1000, 40000)), + expC: 100, // no thinning because pid and no-pid + }, + } { + t.Run(d.name, func(t *testing.T) { + filter, count := makeFilter(d.m) + _ = filter + if d.expC != count { + t.Errorf("expected count %d, got %d", d.expC, count) + } + }) + } +}