chore(buildkit): Add test cases for builder

Signed-off-by: Ce Gao <cegao@tensorchord.ai>

Makefile

@@ -91,7 +91,7 @@ export GOFLAGS ?= -count=1
 #
 
 # All targets.
-.PHONY: lint test build container push addlicense debug debug-local build-local
+.PHONY: lint test build container push addlicense debug debug-local build-local generate test
 
 build: build-local
 
@@ -111,6 +111,10 @@ build-local:
 	    $(CMD_DIR)/$${target};                                                         \
 	done
 
+generate:
+	@mockgen -source pkg/buildkitd/buildkitd.go -destination pkg/buildkitd/mock/mock.go -package mock
+	@mockgen -source pkg/lang/frontend/starlark/interpreter.go -destination pkg/lang/frontend/starlark/mock/mock.go -package mock
+
 # It is used by vscode to attach into the process.
 debug-local:
 	@for target in $(TARGETS); do                                                      \
@@ -123,6 +127,10 @@ debug-local:
 addlicense:
 	addlicense -c "The MIDI Authors" **/*.go **/**/*.go
 
+test: generate
+	@go test -race -coverprofile=coverage.out ./...
+	@go tool cover -func coverage.out | tail -n 1 | awk '{ print "Total coverage: " $$3 }'
+
 .PHONY: clean
 clean:
 	@-rm -vrf ${OUTPUT_DIR}

cmd/midi/bootstrap.go

@@ -73,7 +73,10 @@ func bootstrap(clicontext *cli.Context) error {
 
 	if buildkit {
 		logrus.Debug("bootstrap the buildkitd container")
-		bkClient := buildkitd.NewClient()
+		bkClient, err := buildkitd.NewClient(clicontext.Context)
+		if err != nil {
+			return errors.Wrap(err, "failed to create buildkit client")
+		}
 		defer bkClient.Close()
 		addr, err := bkClient.Bootstrap(clicontext.Context)
 		if err != nil {
@@ -95,7 +98,7 @@ func insertZSHCompleteEntry() error {
 		return errors.Wrapf(err, "failed to check if %s exists", dirPath)
 	}
 	if !dirPathExists {
-		log.L.Warn("Warning: unable to enable zsh-completion: %s does not exist\n", dirPath)
+		log.L.Warnf("Warning: unable to enable zsh-completion: %s does not exist", dirPath)
 		return nil // zsh-completion isn't available, silently fail.
 	}
 

cmd/midi/build.go

@@ -15,15 +15,12 @@
 package main
 
 import (
-	"fmt"
 	"path/filepath"
 
 	"github.com/cockroachdb/errors"
-	"github.com/spf13/viper"
 	cli "github.com/urfave/cli/v2"
 
 	"github.com/tensorchord/MIDI/pkg/builder"
-	"github.com/tensorchord/MIDI/pkg/flag"
 	"github.com/tensorchord/MIDI/pkg/home"
 )
 
@@ -62,7 +59,9 @@ func build(clicontext *cli.Context) error {
 
 	tag := clicontext.String("tag")
 
-	builder := builder.New(fmt.Sprintf("docker-container://%s",
-		viper.GetString(flag.FlagBuildkitdContainer)), config, path, tag)
+	builder, err := builder.New(clicontext.Context, config, path, tag)
+	if err != nil {
+		return errors.Wrap(err, "failed to create the builder")
+	}
 	return builder.Build(clicontext.Context)
 }

cmd/midi/up.go

@@ -15,18 +15,15 @@
 package main
 
 import (
-	"fmt"
 	"path/filepath"
 	"time"
 
 	"github.com/cockroachdb/errors"
 	"github.com/sirupsen/logrus"
-	"github.com/spf13/viper"
 	cli "github.com/urfave/cli/v2"
 
 	"github.com/tensorchord/MIDI/pkg/builder"
 	"github.com/tensorchord/MIDI/pkg/docker"
-	"github.com/tensorchord/MIDI/pkg/flag"
 	"github.com/tensorchord/MIDI/pkg/home"
 	"github.com/tensorchord/MIDI/pkg/ssh"
 )
@@ -82,8 +79,11 @@ func up(clicontext *cli.Context) error {
 
 	tag := clicontext.String("tag")
 
-	builder := builder.New(fmt.Sprintf("docker-container://%s",
-		viper.GetString(flag.FlagBuildkitdContainer)), config, path, tag)
+	builder, err := builder.New(clicontext.Context, config, path, tag)
+	if err != nil {
+		return errors.Wrap(err, "failed to create the builder")
+	}
+
 	if err := builder.Build(clicontext.Context); err != nil {
 		return err
 	}

go.mod

@@ -10,8 +10,11 @@ require (
 	github.com/docker/docker v20.10.7+incompatible
 	github.com/fatih/color v1.13.0
 	github.com/gliderlabs/ssh v0.3.3
+	github.com/golang/mock v1.6.0
 	github.com/google/uuid v1.3.0
 	github.com/moby/buildkit v0.10.1
+	github.com/onsi/ginkgo/v2 v2.1.4
+	github.com/onsi/gomega v1.19.0
 	github.com/opencontainers/go-digest v1.0.0
 	github.com/pkg/sftp v1.13.4
 	github.com/sirupsen/logrus v1.8.1
@@ -20,7 +23,7 @@ require (
 	go.starlark.net v0.0.0-20220328144851-d1966c6b9fcd
 	golang.org/x/crypto v0.0.0-20211202192323-5770296d904e
 	golang.org/x/sync v0.0.0-20210220032951-036812b2e83c
-	golang.org/x/term v0.0.0-20210615171337-6886f2dfbf5b
+	golang.org/x/term v0.0.0-20210927222741-03fcf44c2211
 )
 
 require (
@@ -79,8 +82,8 @@ require (
 	go.opentelemetry.io/otel/sdk v1.4.1 // indirect
 	go.opentelemetry.io/otel/trace v1.4.1 // indirect
 	go.opentelemetry.io/proto/otlp v0.12.0 // indirect
-	golang.org/x/net v0.0.0-20211216030914-fe4d6282115f // indirect
-	golang.org/x/sys v0.0.0-20220209214540-3681064d5158 // indirect
+	golang.org/x/net v0.0.0-20220225172249-27dd8689420f // indirect
+	golang.org/x/sys v0.0.0-20220319134239-a9b59b0215f8 // indirect
 	golang.org/x/text v0.3.7 // indirect
 	golang.org/x/time v0.0.0-20210723032227-1f47c861a9ac // indirect
 	google.golang.org/genproto v0.0.0-20211208223120-3a66f561d7aa // indirect

go.sum

@@ -216,6 +216,8 @@ github.com/golang/mock v1.4.1/go.mod h1:UOMv5ysSaYNkG+OFQykRIcU/QvvxJf3p21QfJ2Bt
 github.com/golang/mock v1.4.3/go.mod h1:UOMv5ysSaYNkG+OFQykRIcU/QvvxJf3p21QfJ2Bt3cw=
 github.com/golang/mock v1.4.4/go.mod h1:l3mdAwkq5BuhzHwde/uurv3sEJeZMXNpwsxVWU71h+4=
 github.com/golang/mock v1.5.0/go.mod h1:CWnOUgYIOo4TcNZ0wHX3YZCqsaM1I1Jvs6v3mP3KVu8=
+github.com/golang/mock v1.6.0 h1:ErTB+efbowRARo13NNdxyJji2egdxLGQhRaY+DUumQc=
+github.com/golang/mock v1.6.0/go.mod h1:p6yTPP+5HYm5mzsMV8JkE6ZKdX+/wYM6Hr+LicevLPs=
 github.com/golang/protobuf v1.2.0/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.1/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
 github.com/golang/protobuf v1.3.2/go.mod h1:6lQm79b+lXiMfvg/cZm0SGofjICqVBUtrP5yJMmIC1U=
@@ -398,8 +400,12 @@ github.com/onsi/ginkgo v1.6.0/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+W
 github.com/onsi/ginkgo v1.10.3/go.mod h1:lLunBs/Ym6LB5Z9jYTR76FiuTmxDTDusOGeTQH+WWjE=
 github.com/onsi/ginkgo v1.12.1/go.mod h1:zj2OWP4+oCPe1qIXoGWkgMRwljMUYCdkwsT2108oapk=
 github.com/onsi/ginkgo v1.13.0/go.mod h1:+REjRxOmWfHCjfv9TTWB1jD1Frx4XydAD3zm1lskyM0=
+github.com/onsi/ginkgo/v2 v2.1.4 h1:GNapqRSid3zijZ9H77KrgVG4/8KqiyRsxcSxe+7ApXY=
+github.com/onsi/ginkgo/v2 v2.1.4/go.mod h1:um6tUpWM/cxCK3/FK8BXqEiUMUwRgSM4JXG47RKZmLU=
 github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY=
 github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo=
+github.com/onsi/gomega v1.19.0 h1:4ieX6qQjPP/BfC3mpsAtIGGlxTWPeA3Inl/7DtXw1tw=
+github.com/onsi/gomega v1.19.0/go.mod h1:LY+I3pBVzYsTBU1AnDwOSxaYi9WoWiqgwooUqq9yPro=
 github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U=
 github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM=
 github.com/opencontainers/image-spec v1.0.3-0.20211202183452-c5a74bcca799 h1:rc3tiVYb5z54aKaDfakKn0dDjIyPpTtszkjuMzyt7ec=
@@ -636,8 +642,8 @@ golang.org/x/net v0.0.0-20210226172049-e18ecbb05110/go.mod h1:m0MpNAwzfU5UDzcl9v
 golang.org/x/net v0.0.0-20210316092652-d523dce5a7f4/go.mod h1:RBQZq4jEuRlivfhVLdyRGr576XBO4/greRjx4P4O3yc=
 golang.org/x/net v0.0.0-20210405180319-a5a99cb37ef4/go.mod h1:p54w0d4576C0XHj96bSt6lcn1PtDYWL6XObtHCRCNQM=
 golang.org/x/net v0.0.0-20211008194852-3b03d305991f/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
-golang.org/x/net v0.0.0-20211216030914-fe4d6282115f h1:hEYJvxw1lSnWIl8X9ofsYMklzaDs90JI2az5YMd4fPM=
-golang.org/x/net v0.0.0-20211216030914-fe4d6282115f/go.mod h1:9nx3DQGgdP8bBQD5qxJ1jj9UTztislL4KSBs9R2vV5Y=
+golang.org/x/net v0.0.0-20220225172249-27dd8689420f h1:oA4XRj0qtSt8Yo1Zms0CUlsT3KG69V2UGQWPBxujDmc=
+golang.org/x/net v0.0.0-20220225172249-27dd8689420f/go.mod h1:CfG3xpIq0wQ8r1q4Su4UZFWDARRcnwPjda9FqA0JpMk=
 golang.org/x/oauth2 v0.0.0-20180821212333-d2e6202438be/go.mod h1:N/0e6XlmueqKjAGxoOufVs8QHGRruUQn6yWY3a++T0U=
 golang.org/x/oauth2 v0.0.0-20190226205417-e64efc72b421/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
 golang.org/x/oauth2 v0.0.0-20190604053449-0f29369cfe45/go.mod h1:gOpvHmFTYa4IltrdGE7lF6nIHvwfUNPOp7c8zoXwtLw=
@@ -726,11 +732,12 @@ golang.org/x/sys v0.0.0-20210630005230-0f9fa26af87c/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.0.0-20210927094055-39ccf1dd6fa6/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20211007075335-d3039528d8ac/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20211025201205-69cdffdb9359/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
-golang.org/x/sys v0.0.0-20220209214540-3681064d5158 h1:rm+CHSpPEEW2IsXUib1ThaHIjuBVZjxNgSKmBLFfD4c=
 golang.org/x/sys v0.0.0-20220209214540-3681064d5158/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.0.0-20220319134239-a9b59b0215f8 h1:OH54vjqzRWmbJ62fjuhxy7AxFFgoHN0/DPc/UrL8cAs=
+golang.org/x/sys v0.0.0-20220319134239-a9b59b0215f8/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
-golang.org/x/term v0.0.0-20210615171337-6886f2dfbf5b h1:9zKuko04nR4gjZ4+DNjHqRlAJqbJETHwiNKDqTfOjfE=
-golang.org/x/term v0.0.0-20210615171337-6886f2dfbf5b/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
+golang.org/x/term v0.0.0-20210927222741-03fcf44c2211 h1:JGgROgKl9N8DuW20oFS5gxc+lE67/N3FcwmBPMe7ArY=
+golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 golang.org/x/text v0.0.0-20170915032832-14c0d48ead0c/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.1-0.20180807135948-17ff2d5776d2/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
@@ -801,6 +808,7 @@ golang.org/x/tools v0.0.0-20201208233053-a543418bbed2/go.mod h1:emZCQorbCU4vsT4f
 golang.org/x/tools v0.0.0-20210105154028-b0ab187a4818/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/tools v0.0.0-20210106214847-113979e3529a/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA=
 golang.org/x/tools v0.1.0/go.mod h1:xkSsbof2nBLbhDlRMhhhyNLN/zl3eTqcnHD5viDpcZ0=
+golang.org/x/tools v0.1.1/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/tools v0.1.3/go.mod h1:o0xws9oXOQQZyjljx8fwUC0k7L1pTE6eaCbjGeHmOkk=
 golang.org/x/xerrors v0.0.0-20190717185122-a985d3407aa7/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=

pkg/builder/builder.go

@@ -23,12 +23,14 @@ import (
 	"github.com/moby/buildkit/client"
 	"github.com/moby/buildkit/util/progress/progresswriter"
 	"github.com/sirupsen/logrus"
+	"golang.org/x/sync/errgroup"
+
+	"github.com/tensorchord/MIDI/pkg/buildkitd"
 	"github.com/tensorchord/MIDI/pkg/docker"
 	"github.com/tensorchord/MIDI/pkg/flag"
 	"github.com/tensorchord/MIDI/pkg/home"
 	"github.com/tensorchord/MIDI/pkg/lang/frontend/starlark"
 	"github.com/tensorchord/MIDI/pkg/lang/ir"
-	"golang.org/x/sync/errgroup"
 )
 
 type Builder interface {
@@ -37,28 +39,36 @@ type Builder interface {
 }
 
 type generalBuilder struct {
-	buildkitdSocket  string
 	manifestFilePath string
 	configFilePath   string
 	progressMode     string
 	tag              string
 
 	logger *logrus.Entry
+	starlark.Interpreter
+	buildkitd.Client
 }
 
-func New(buildkitdSocket, configFilePath, manifestFilePath, tag string) Builder {
-	return &generalBuilder{
-		buildkitdSocket:  buildkitdSocket,
+func New(ctx context.Context, configFilePath, manifestFilePath, tag string) (Builder, error) {
+	b := &generalBuilder{
 		manifestFilePath: manifestFilePath,
 		configFilePath:   configFilePath,
 		// TODO(gaocegege): Support other mode?
 		progressMode: "auto",
 		tag:          tag,
 		logger: logrus.WithFields(logrus.Fields{
-			"tag":       tag,
-			"buildkitd": buildkitdSocket,
+			"tag": tag,
 		}),
 	}
+
+	cli, err := buildkitd.NewClient(ctx)
+	if err != nil {
+		return nil, errors.Wrap(err, "failed to create buildkit client")
+	}
+	b.Client = cli
+
+	b.Interpreter = starlark.NewInterpreter()
+	return b, nil
 }
 
 // GPUEnabled returns true if cuda is enabled.
@@ -68,22 +78,15 @@ func (b generalBuilder) GPUEnabled() bool {
 }
 
 func (b generalBuilder) Build(ctx context.Context) error {
-	interpreter := starlark.NewInterpreter()
 	// Evaluate config first.
-	if _, err := interpreter.ExecFile(b.configFilePath); err != nil {
+	if _, err := b.ExecFile(b.configFilePath); err != nil {
 		return err
 	}
 
-	if _, err := interpreter.ExecFile(b.manifestFilePath); err != nil {
+	if _, err := b.ExecFile(b.manifestFilePath); err != nil {
 		return err
 	}
 
-	bkClient, err := client.New(ctx, b.buildkitdSocket, client.WithFailFast())
-	if err != nil {
-		return errors.Wrap(err, "failed to new buildkitd client")
-	}
-	defer bkClient.Close()
-
 	def, err := ir.Compile(ctx)
 	if err != nil {
 		return errors.Wrap(err, "failed to compile build.MIDI")
@@ -106,7 +109,7 @@ func (b generalBuilder) Build(ctx context.Context) error {
 			return err
 		}
 		b.logger.Debug("building image in ", wd)
-		_, err = bkClient.Solve(ctx, def, client.SolveOpt{
+		_, err = b.Solve(ctx, def, client.SolveOpt{
 			Exports: []client.ExportEntry{
 				{
 					Type: client.ExporterDocker,

pkg/builder/builder_suite_test.go

@@ -0,0 +1,13 @@
+package builder
+
+import (
+	"testing"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+func TestOci(t *testing.T) {
+	RegisterFailHandler(Fail)
+	RunSpecs(t, "Builder Suite")
+}

pkg/builder/builder_test.go

@@ -0,0 +1,30 @@
+package builder
+
+import (
+	"context"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+	"github.com/spf13/viper"
+	"github.com/tensorchord/MIDI/pkg/flag"
+)
+
+var _ = Describe("Builder", func() {
+	Describe("building", Label("buildkitd"), func() {
+		var buildkitdSocket, configFilePath, manifestFilePath, tag string
+		BeforeEach(func() {
+			buildkitdSocket = "docker-container://midi_buildkitd"
+			configFilePath = "testdata/config.MIDI"
+			manifestFilePath = "testdata/build.MIDI"
+			tag = "midi-dev:test"
+		})
+		When("getting the wrong builtkitd address", func() {
+			buildkitdSocket = "wrong"
+			viper.Set(flag.FlagBuildkitdContainer, buildkitdSocket)
+			It("should return an error", func() {
+				_, err := New(context.TODO(), configFilePath, manifestFilePath, tag)
+				Expect(err).To(HaveOccurred())
+			})
+		})
+	})
+})