1 files changed, 20 insertions, 5 deletions
diff --git a/pkg/test/dockerutil/container.go b/pkg/test/dockerutil/container.go
index b59503188..5a2157951 100644
--- a/pkg/test/dockerutil/container.go
+++ b/pkg/test/dockerutil/container.go
@@ -360,13 +360,18 @@ func (c *Container) SandboxPid(ctx context.Context) (int, error) {
 }
 
 // FindIP returns the IP address of the container.
-func (c *Container) FindIP(ctx context.Context) (net.IP, error) {
+func (c *Container) FindIP(ctx context.Context, ipv6 bool) (net.IP, error) {
 	resp, err := c.client.ContainerInspect(ctx, c.id)
 	if err != nil {
 		return nil, err
 	}
 
-	ip := net.ParseIP(resp.NetworkSettings.DefaultNetworkSettings.IPAddress)
+	var ip net.IP
+	if ipv6 {
+		ip = net.ParseIP(resp.NetworkSettings.DefaultNetworkSettings.GlobalIPv6Address)
+	} else {
+		ip = net.ParseIP(resp.NetworkSettings.DefaultNetworkSettings.IPAddress)
+	}
 	if ip == nil {
 		return net.IP{}, fmt.Errorf("invalid IP: %q", ip)
 	}
@@ -449,15 +454,19 @@ func (c *Container) Wait(ctx context.Context) error {
 
 // WaitTimeout waits for the container to exit with a timeout.
 func (c *Container) WaitTimeout(ctx context.Context, timeout time.Duration) error {
-	timeoutChan := time.After(timeout)
+	ctx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
 	statusChan, errChan := c.client.ContainerWait(ctx, c.id, container.WaitConditionNotRunning)
 	select {
+	case <-ctx.Done():
+		if ctx.Err() == context.DeadlineExceeded {
+			return fmt.Errorf("container %s timed out after %v seconds", c.Name, timeout.Seconds())
+		}
+		return nil
 	case err := <-errChan:
 		return err
 	case <-statusChan:
 		return nil
-	case <-timeoutChan:
-		return fmt.Errorf("container %s timed out after %v seconds", c.Name, timeout.Seconds())
 	}
 }
 
@@ -482,6 +491,12 @@ func (c *Container) WaitForOutputSubmatch(ctx context.Context, pattern string, t
 	}
 
 	for exp := time.Now().Add(timeout); time.Now().Before(exp); {
+		select {
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		default:
+		}
+
 		c.streams.Conn.SetDeadline(time.Now().Add(50 * time.Millisecond))
 		_, err := stdcopy.StdCopy(&c.streamBuf, &c.streamBuf, c.streams.Reader)