From bc6f2fd49d525c139ea199648ba2fee1c996d71b Mon Sep 17 00:00:00 2001
From: Pasha Sviderski <me@psviderski.name>
Date: Thu, 25 Sep 2025 20:16:56 +1000
Subject: [PATCH] chore: fix docker client CreateContainerWithImagePull, add
 WaitPortPublished

---
 internal/docker/container.go | 48 +++++++++++++++++++++++++++---------
 1 file changed, 37 insertions(+), 11 deletions(-)

diff --git a/internal/docker/container.go b/internal/docker/container.go
index 962d2486..202a3bc2 100644
--- a/internal/docker/container.go
+++ b/internal/docker/container.go
@@ -2,26 +2,23 @@ package docker
 
 import (
 	"context"
+	"errors"
 	"fmt"
+	"time"
 
 	"github.com/docker/docker/api/types/container"
 	"github.com/docker/docker/api/types/image"
 	"github.com/docker/docker/client"
+	"github.com/docker/go-connections/nat"
 )
 
-// CreateContainerWithImagePull creates a Docker container. If the image is missing, it pulls the image first.
+// CreateContainerWithImagePull creates a container. If the image is missing, it pulls the image first.
 func (cli *Client) CreateContainerWithImagePull(
 	ctx context.Context, name string, config *container.Config, hostConfig *container.HostConfig,
 ) (container.CreateResponse, error) {
-	var resp container.CreateResponse
-
-	_, err := cli.ContainerCreate(ctx, config, hostConfig, nil, nil, name)
-	if err == nil {
-		return resp, nil
-	}
-
-	if !client.IsErrNotFound(err) {
-		return resp, fmt.Errorf("create container: %w", err)
+	resp, err := cli.ContainerCreate(ctx, config, hostConfig, nil, nil, name)
+	if err == nil || !client.IsErrNotFound(err) {
+		return resp, err
 	}
 
 	pullCh, err := cli.PullImage(ctx, config.Image, image.PullOptions{})
@@ -29,6 +26,8 @@ func (cli *Client) CreateContainerWithImagePull(
 		return resp, fmt.Errorf("pull image: %w", err)
 	}
 
+	// Drain the pull channel until it's closed (image is fully pulled) or an error occurs.
+	// If the context is canceled during the pull, the channel will receive a context cancellation error.
 	for msg := range pullCh {
 		if msg.Err != nil {
 			return resp, fmt.Errorf("pull image: %w", msg.Err)
@@ -37,8 +36,35 @@ func (cli *Client) CreateContainerWithImagePull(
 
 	// Create container again after image pull.
 	if resp, err = cli.ContainerCreate(ctx, config, hostConfig, nil, nil, name); err != nil {
-		return resp, fmt.Errorf("create container: %w", err)
+		return resp, err
 	}
 
 	return resp, nil
 }
+
+// WaitPortPublished waits for a container port to be published on the host which happens asynchronously.
+func (cli *Client) WaitPortPublished(ctx context.Context, containerID string, port nat.Port) ([]nat.PortBinding, error) {
+	ctx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+
+	for {
+		c, err := cli.ContainerInspect(ctx, containerID)
+		if err != nil {
+			return nil, fmt.Errorf("inspect container: %w", err)
+		}
+
+		binding, ok := c.NetworkSettings.Ports[port]
+		if ok && len(binding) > 0 {
+			return binding, nil
+		}
+
+		select {
+		case <-time.After(10 * time.Millisecond):
+		case <-ctx.Done():
+			if errors.Is(ctx.Err(), context.DeadlineExceeded) {
+				return nil, errors.New("timeout")
+			}
+			return nil, ctx.Err()
+		}
+	}
+}