3 files changed, 15 insertions, 5 deletions
diff --git a/test/packetimpact/testbench/connections.go b/test/packetimpact/testbench/connections.go
index 56ac3fa54..28c841612 100644
--- a/test/packetimpact/testbench/connections.go
+++ b/test/packetimpact/testbench/connections.go
@@ -841,6 +841,7 @@ func (conn *UDPIPv4) SendIP(additionalLayers ...Layer) {
 // Expect expects a frame with the UDP layer matching the provided UDP within
 // the timeout specified. If it doesn't arrive in time, an error is returned.
 func (conn *UDPIPv4) Expect(udp UDP, timeout time.Duration) (*UDP, error) {
+	conn.t.Helper()
 	layer, err := (*Connection)(conn).Expect(&udp, timeout)
 	if layer == nil {
 		return nil, err
@@ -852,6 +853,18 @@ func (conn *UDPIPv4) Expect(udp UDP, timeout time.Duration) (*UDP, error) {
 	return gotUDP, err
 }
 
+// ExpectData is a convenient method that expects a Layer and the Layer after
+// it. If it doens't arrive in time, it returns nil.
+func (conn *UDPIPv4) ExpectData(udp UDP, payload Payload, timeout time.Duration) (Layers, error) {
+	conn.t.Helper()
+	expected := make([]Layer, len(conn.layerStates))
+	expected[len(expected)-1] = &udp
+	if payload.length() != 0 {
+		expected = append(expected, &payload)
+	}
+	return (*Connection)(conn).ExpectFrame(expected, timeout)
+}
+
 // Close frees associated resources held by the UDPIPv4 connection.
 func (conn *UDPIPv4) Close() {
 	(*Connection)(conn).Close()
diff --git a/test/packetimpact/testbench/layers.go b/test/packetimpact/testbench/layers.go
index 165f62d3b..49370377d 100644
--- a/test/packetimpact/testbench/layers.go
+++ b/test/packetimpact/testbench/layers.go
@@ -898,10 +898,7 @@ func (l *UDP) match(other Layer) bool {
 }
 
 func (l *UDP) length() int {
-	if l.Length == nil {
-		return header.UDPMinimumSize
-	}
-	return int(*l.Length)
+	return header.UDPMinimumSize
 }
 
 // merge implements Layer.merge.
diff --git a/test/packetimpact/testbench/rawsockets.go b/test/packetimpact/testbench/rawsockets.go
index ff722d4a6..a9ad72b63 100644
--- a/test/packetimpact/testbench/rawsockets.go
+++ b/test/packetimpact/testbench/rawsockets.go
@@ -169,7 +169,7 @@ func NewInjector(t *testing.T) (Injector, error) {
 // Send a raw frame.
 func (i *Injector) Send(b []byte) {
 	if _, err := unix.Write(i.fd, b); err != nil {
-		i.t.Fatalf("can't write: %s", err)
+		i.t.Fatalf("can't write: %s of len %d", err, len(b))
 	}
 }