diff options
Diffstat (limited to 'test/packetimpact/testbench/connections.go')
| -rw-r--r-- | test/packetimpact/testbench/connections.go | 13 |
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() |