diff --git a/pkg/universe.dagger.io/yarn/test/test.cue b/pkg/universe.dagger.io/yarn/test/test.cue
index fa5ed054..dea868f7 100644
--- a/pkg/universe.dagger.io/yarn/test/test.cue
+++ b/pkg/universe.dagger.io/yarn/test/test.cue
@@ -4,6 +4,7 @@ import (
 	"dagger.io/dagger"
 	"dagger.io/dagger/engine"
 
+	"universe.dagger.io/docker"
 	"universe.dagger.io/yarn"
 )
 
@@ -15,10 +16,15 @@ dagger.#Plan & {
 
 	actions: tests: {
 
+		// Configuration for all tests
+		common: {
+			data: inputs.directories.testdata.contents
+		}
+
 		// Run yarn.#Build
 		simple: {
 			build: yarn.#Build & {
-				source: inputs.directories.testdata.contents
+				source: common.data
 			}
 
 			verify: #AssertFile & {
@@ -32,7 +38,7 @@ dagger.#Plan & {
 		customName: {
 			build: yarn.#Build & {
 				name:   "My Build"
-				source: inputs.directories.testdata.contents
+				source: common.data
 			}
 			verify: #AssertFile & {
 				input:    build.output
@@ -40,6 +46,30 @@ dagger.#Plan & {
 				contents: "output\n"
 			}
 		}
+
+		// Run yarn.#Build with a custom docker image
+		customImage: {
+			buildImage: docker.#Build & {
+				steps: [
+					docker.#Pull & {
+						source: "alpine"
+					},
+					docker.#Run & {
+						command: {
+							name: "apk"
+							args: ["add", "yarn", "bash"]
+						}
+					},
+				]
+			}
+
+			image: build.output
+
+			build: yarn.#Build & {
+				source: common.data
+				container: input: buildImage.output
+			}
+		}
 	}
 }
 
diff --git a/pkg/universe.dagger.io/yarn/yarn.cue b/pkg/universe.dagger.io/yarn/yarn.cue
index fc092a38..2cc5bc19 100644
--- a/pkg/universe.dagger.io/yarn/yarn.cue
+++ b/pkg/universe.dagger.io/yarn/yarn.cue
@@ -9,6 +9,7 @@ import (
 
 	"universe.dagger.io/alpine"
 	"universe.dagger.io/bash"
+	"universe.dagger.io/docker"
 )
 
 // Build a Yarn package
@@ -59,7 +60,7 @@ import (
 
 	// Run yarn in a docker container
 	container: bash.#Run & {
-		input: _buildImage.output
+		input: *_buildImage.output | docker.#Image
 
 		// FIXME: move shell script to its own file
 		script: contents: #"""