Handle context cancellation during builds (#539)

`buildx.Build` doesn't terminate if context is canceled, so this PR
rearranges things such that we can wait for the build or context to
finish.

Fixes #533.

provider/internal/client_test.go

@@ -17,15 +17,23 @@ package internal
 import (
 	"bytes"
 	"context"
+	"errors"
 	"io"
 	"log/slog"
 	"os"
 	"path/filepath"
 	"testing"
 
+	buildx "github.com/docker/buildx/build"
+	"github.com/docker/buildx/builder"
+	"github.com/docker/buildx/util/confutil"
+	"github.com/docker/buildx/util/dockerutil"
+	"github.com/docker/buildx/util/progress"
 	"github.com/docker/docker/api/types/registry"
+	"github.com/moby/buildkit/client"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/mock/gomock"
 )
 
 func TestAuth(t *testing.T) {
@@ -432,6 +440,35 @@ func TestBuildExecError(t *testing.T) {
 	}
 }
 
+func TestBuildCancelation(t *testing.T) {
+	t.Parallel()
+	cli := testcli(t, true)
+
+	ctrl := gomock.NewController(t)
+
+	ctx, cancel := context.WithCancel(context.Background())
+
+	b := NewMockBuilder(ctrl)
+	b.EXPECT().Build(
+		gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(), gomock.Any(),
+	).DoAndReturn(func(
+		_ context.Context,
+		_ []builder.Node,
+		_ map[string]buildx.Options,
+		_ *dockerutil.Client,
+		_ *confutil.Config,
+		_ progress.Writer,
+	) (map[string]*client.SolveResponse, error) {
+		cancel()
+		return nil, errors.New("cancel wasn't respected")
+	})
+	cli.builder = b
+
+	resp, err := cli.Build(ctx, &build{})
+	assert.ErrorIs(t, err, context.Canceled)
+	assert.Nil(t, resp)
+}
+
 // testcli returns a new standalone CLI instance. Set ping to true if a live
 // daemon is required -- the test will be skipped if the daemon is not available.
 func testcli(t *testing.T, ping bool, auths ...Registry) *cli {