Add UDP send/recv packetimpact tests.

Fixes #2654

PiperOrigin-RevId: 310642216

1 files changed, 13 insertions, 0 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()