feat: allow saving specific platforms for an image

Author: LaurentGoderre

docker.go

@@ -1786,6 +1786,22 @@ func (p *DockerProvider) ListImages(ctx context.Context) ([]ImageInfo, error) {
 
 // SaveImages exports a list of images as an uncompressed tar
 func (p *DockerProvider) SaveImages(ctx context.Context, output string, images ...string) error {
+	return p.SaveImagesWithOps(ctx, output, images)
+}
+
+// SaveImagesWithOpts exports a list of images as an uncompressed tar, passiong options to thw provider
+func (p *DockerProvider) SaveImagesWithOps(ctx context.Context, output string, images []string, opts ...SaveImageOption) error {
+	saveOpts := saveOptions{}
+	dockerOpts := []client.ImageSaveOption{}
+
+	for _, o := range opts {
+		o(&saveOpts)
+	}
+
+	if len(saveOpts.Platforms) > 0 {
+		dockerOpts = append(dockerOpts, client.ImageSaveWithPlatforms(saveOpts.Platforms...))
+	}
+
 	outputFile, err := os.Create(output)
 	if err != nil {
 		return fmt.Errorf("opening output file %w", err)
@@ -1794,7 +1810,7 @@ func (p *DockerProvider) SaveImages(ctx context.Context, output string, images .
 		_ = outputFile.Close()
 	}()
 
-	imageReader, err := p.client.ImageSave(ctx, images)
+	imageReader, err := p.client.ImageSave(ctx, images, dockerOpts...)
 	if err != nil {
 		return fmt.Errorf("saving images %w", err)
 	}

image.go

@@ -2,6 +2,9 @@ package testcontainers
 
 import (
 	"context"
+
+	"github.com/containerd/platforms"
+	specs "github.com/opencontainers/image-spec/specs-go/v1"
 )
 
 // ImageInfo represents summary information of an image
@@ -10,9 +13,31 @@ type ImageInfo struct {
 	Name string
 }
 
+type saveOptions struct {
+	Platforms []specs.Platform
+}
+
+type SaveImageOption func(*saveOptions) error
+
+// SaveImageWithPlatforms allows specifying which platform(s) to save
+func SaveImageWithPlatforms(plaforms ...string) SaveImageOption {
+	return func(opts *saveOptions) error {
+		for _, p := range plaforms {
+			platform, err := platforms.Parse(p)
+			if err != nil {
+				return err
+			}
+			opts.Platforms = append(opts.Platforms, platform)
+		}
+
+		return nil
+	}
+}
+
 // ImageProvider allows manipulating images
 type ImageProvider interface {
 	ListImages(context.Context) ([]ImageInfo, error)
 	SaveImages(context.Context, string, ...string) error
+	SaveImagesWithOps(context.Context, string, []string, ...SaveImageOption) error
 	PullImage(context.Context, string) error
 }

image_test.go

@@ -7,7 +7,6 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/require"
-
 	"github.com/testcontainers/testcontainers-go/internal/core"
 )
 
@@ -18,7 +17,8 @@ func TestImageList(t *testing.T) {
 	require.NoErrorf(t, err, "failed to get provider")
 
 	defer func() {
-		_ = provider.Close()
+		err = provider.Close()
+		require.NoError(t, err)
 	}()
 
 	req := ContainerRequest{
@@ -51,7 +51,8 @@ func TestSaveImages(t *testing.T) {
 	require.NoErrorf(t, err, "failed to get provider")
 
 	defer func() {
-		_ = provider.Close()
+		err = provider.Close()
+		require.NoError(t, err)
 	}()
 
 	req := ContainerRequest{
@@ -71,3 +72,35 @@ func TestSaveImages(t *testing.T) {
 
 	require.NotZerof(t, info.Size(), "output file is empty")
 }
+
+func TestSaveImagesWithOpts(t *testing.T) {
+	t.Setenv("DOCKER_HOST", core.MustExtractDockerHost(context.Background()))
+
+	provider, err := ProviderDocker.GetProvider()
+	require.NoErrorf(t, err, "failed to get provider")
+
+	defer func() {
+		err = provider.Close()
+		require.NoError(t, err)
+	}()
+
+	req := ContainerRequest{
+		Image:         "redis:latest",
+		ImagePlatform: "linux/amd64",
+	}
+
+	ctr, err := provider.CreateContainer(context.Background(), req)
+	CleanupContainer(t, ctr)
+	require.NoErrorf(t, err, "creating test container")
+
+	output := filepath.Join(t.TempDir(), "images.tar")
+	err = provider.SaveImagesWithOps(
+		context.Background(), output, []string{req.Image}, SaveImageWithPlatforms("linux/amd64"),
+	)
+	require.NoErrorf(t, err, "saving image %q", req.Image)
+
+	info, err := os.Stat(output)
+	require.NoError(t, err)
+
+	require.NotZerof(t, info.Size(), "output file is empty")
+}

modules/k3s/k3s.go

@@ -195,7 +195,13 @@ func unmarshal(bytes []byte) (*KubeConfigValue, error) {
 }
 
 // LoadImages loads images into the k3s container.
+type LoadImageOption = testcontainers.SaveImageOption
+
 func (c *K3sContainer) LoadImages(ctx context.Context, images ...string) error {
+	return c.LoadImagesWithOpts(ctx, images)
+}
+
+func (c *K3sContainer) LoadImagesWithOpts(ctx context.Context, images []string, opts ...LoadImageOption) error {
 	provider, err := testcontainers.ProviderDocker.GetProvider()
 	if err != nil {
 		return fmt.Errorf("getting docker provider %w", err)
@@ -210,7 +216,7 @@ func (c *K3sContainer) LoadImages(ctx context.Context, images ...string) error {
 		_ = os.Remove(imagesTar.Name())
 	}()
 
-	err = provider.SaveImages(context.Background(), imagesTar.Name(), images...)
+	err = provider.SaveImagesWithOps(context.Background(), imagesTar.Name(), images, opts...)
 	if err != nil {
 		return fmt.Errorf("saving images %w", err)
 	}
@@ -221,7 +227,7 @@ func (c *K3sContainer) LoadImages(ctx context.Context, images ...string) error {
 		return fmt.Errorf("copying image to container %w", err)
 	}
 
-	_, _, err = c.Exec(ctx, []string{"ctr", "-n=k8s.io", "images", "import", containerPath})
+	_, _, err = c.Exec(ctx, []string{"ctr", "-n=k8s.io", "images", "import", "--all-platforms", containerPath})
 	if err != nil {
 		return fmt.Errorf("importing image %w", err)
 	}