1 files changed, 12 insertions, 0 deletions
diff --git a/test/packetimpact/testbench/connections.go b/test/packetimpact/testbench/connections.go
index 79c0ccf5c..2b8e2f005 100644
--- a/test/packetimpact/testbench/connections.go
+++ b/test/packetimpact/testbench/connections.go
@@ -228,6 +228,12 @@ func (conn *TCPIPv4) RecvFrame(timeout time.Duration) Layers {
 	return nil
 }
 
+// Drain drains the sniffer's receive buffer by receiving packets until there's
+// nothing else to receive.
+func (conn *TCPIPv4) Drain() {
+	conn.sniffer.Drain()
+}
+
 // Expect a packet that matches the provided tcp within the timeout specified.
 // If it doesn't arrive in time, it returns nil.
 func (conn *TCPIPv4) Expect(tcp TCP, timeout time.Duration) (*TCP, error) {
@@ -423,6 +429,12 @@ func (conn *UDPIPv4) Recv(timeout time.Duration) *UDP {
 	return nil
 }
 
+// Drain drains the sniffer's receive buffer by receiving packets until there's
+// nothing else to receive.
+func (conn *UDPIPv4) Drain() {
+	conn.sniffer.Drain()
+}
+
 // Expect a packet that matches the provided udp within the timeout specified.
 // If it doesn't arrive in time, the test fails.
 func (conn *UDPIPv4) Expect(udp UDP, timeout time.Duration) (*UDP, error) {