tests: connection subset testing

Utility functions to create fake sets of connections for testing, and
then exercising the subset filtering code to check that quantities
come out as expected.

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)
+			}
+		})
+	}
+}