4 files changed, 14 insertions, 13 deletions
diff --git a/test/packetimpact/testbench/dut.go b/test/packetimpact/testbench/dut.go
index 269e163bb..0cac0bf1b 100644
--- a/test/packetimpact/testbench/dut.go
+++ b/test/packetimpact/testbench/dut.go
@@ -498,8 +498,8 @@ func (dut *DUT) PollOne(t *testing.T, fd int32, events int16, timeout time.Durat
 	if readyFd := pfds[0].Fd; readyFd != fd {
 		t.Fatalf("Poll returned an fd %d that was not requested (%d)", readyFd, fd)
 	}
-	if got, want := pfds[0].Revents, int16(events); got&want == 0 {
-		t.Fatalf("Poll returned no events in our interest, got: %#b, want: %#b", got, want)
+	if got, want := pfds[0].Revents, int16(events); got&want != want {
+		t.Fatalf("Poll returned events does not include all of the interested events, got: %#b, want: %#b", got, want)
 	}
 }
 
diff --git a/test/packetimpact/testbench/layers.go b/test/packetimpact/testbench/layers.go
index 078f500e4..2644b3248 100644
--- a/test/packetimpact/testbench/layers.go
+++ b/test/packetimpact/testbench/layers.go
@@ -357,7 +357,7 @@ func (l *IPv4) ToBytes() ([]byte, error) {
 		case *ICMPv4:
 			fields.Protocol = uint8(header.ICMPv4ProtocolNumber)
 		default:
-			// TODO(b/150301488): Support more protocols as needed.
+			// We can add support for more protocols as needed.
 			return nil, fmt.Errorf("ipv4 header's next layer is unrecognized: %#v", n)
 		}
 	}
diff --git a/test/packetimpact/testbench/rawsockets.go b/test/packetimpact/testbench/rawsockets.go
index feeb0888a..6d95c033d 100644
--- a/test/packetimpact/testbench/rawsockets.go
+++ b/test/packetimpact/testbench/rawsockets.go
@@ -17,7 +17,6 @@ package testbench
 import (
 	"encoding/binary"
 	"fmt"
-	"math"
 	"net"
 	"testing"
 	"time"
@@ -81,19 +80,20 @@ func (s *Sniffer) Recv(t *testing.T, timeout time.Duration) []byte {
 
 	deadline := time.Now().Add(timeout)
 	for {
-		timeout = deadline.Sub(time.Now())
+		timeout = time.Until(deadline)
 		if timeout <= 0 {
 			return nil
 		}
-		whole, frac := math.Modf(timeout.Seconds())
-		tv := unix.Timeval{
-			Sec:  int64(whole),
-			Usec: int64(frac * float64(time.Second/time.Microsecond)),
+		usec := timeout.Microseconds()
+		if usec == 0 {
+			// Timeout is less than a microsecond; set usec to 1 to avoid
+			// blocking indefinitely.
+			usec = 1
 		}
-		// The following should never happen, but having this guard here is better
-		// than blocking indefinitely in the future.
-		if tv.Sec == 0 && tv.Usec == 0 {
-			t.Fatal("setting SO_RCVTIMEO to 0 means blocking indefinitely")
+		const microsInOne = 1e6
+		tv := unix.Timeval{
+			Sec:  usec / microsInOne,
+			Usec: usec % microsInOne,
 		}
 		if err := unix.SetsockoptTimeval(s.fd, unix.SOL_SOCKET, unix.SO_RCVTIMEO, &tv); err != nil {
 			t.Fatalf("can't setsockopt SO_RCVTIMEO: %s", err)
diff --git a/test/packetimpact/testbench/testbench.go b/test/packetimpact/testbench/testbench.go
index d0efbcc08..caa389780 100644
--- a/test/packetimpact/testbench/testbench.go
+++ b/test/packetimpact/testbench/testbench.go
@@ -132,6 +132,7 @@ func registerFlags(fs *flag.FlagSet) {
 // Initialize initializes the testbench, it parse the flags and sets up the
 // pool of test networks for testbench's later use.
 func Initialize(fs *flag.FlagSet) {
+	testing.Init()
 	registerFlags(fs)
 	flag.Parse()
 	if err := loadDUTInfos(); err != nil {