[System Probe] Update unit test to reflect UDP connection direction c…

…hanges (#5314)
DataDog · Apr 24, 2020 · a049687 · a049687
1 parent 5546f52
commit a049687
Showing 1 changed file with 13 additions and 3 deletions.
diff --git a/pkg/ebpf/tracer_direction_test.go b/pkg/ebpf/tracer_direction_test.go
@@ -28,11 +28,21 @@ func TestOutgoingTCPDirection(t *testing.T) {
 	assert.Equal(t, OUTGOING, connStat.Direction)
 }
 
-func TestUDPConnectionDirection(t *testing.T) {
+func TestIncomingUDPConnectionDirection(t *testing.T) {
 	tr := &Tracer{
-		portMapping: NewPortMapping(NewDefaultConfig().ProcRoot, NewDefaultConfig()),
+		udpPortMapping: NewPortMapping(NewDefaultConfig().ProcRoot, NewDefaultConfig()),
 	}
+	tr.udpPortMapping.AddMapping(5323)
 	connStat := CreateConnectionStat("10.2.25.1", "38.122.226.210", 5323, 8125, UDP)
 	connStat.Direction = tr.determineConnectionDirection(&connStat)
-	assert.Equal(t, NONE, connStat.Direction)
+	assert.Equal(t, INCOMING, connStat.Direction)
+}
+
+func TestOutgoingUDPConnectionDirection(t *testing.T) {
+	tr := &Tracer{
+		udpPortMapping: NewPortMapping(NewDefaultConfig().ProcRoot, NewDefaultConfig()),
+	}
+	connStat := CreateConnectionStat("10.2.25.1", "38.122.226.210", 5323, 8125, UDP)
+	connStat.Direction = tr.determineConnectionDirection(&connStat)
+	assert.Equal(t, OUTGOING, connStat.Direction)
 }