Fix `down` with `--rmi`

Signed-off-by: Ulysses Souza <[email protected]>

pkg/api/labels.go

@@ -47,6 +47,8 @@ const (
 	OneoffLabel = "com.docker.compose.oneoff"
 	// SlugLabel stores unique slug used for one-off container identity
 	SlugLabel = "com.docker.compose.slug"
+	// ImageNameLabel stores the content of the image section in the compose file
+	ImageNameLabel = "com.docker.compose.image_name"
 	// ImageDigestLabel stores digest of the container image used to run service
 	ImageDigestLabel = "com.docker.compose.image"
 	// DependenciesLabel stores service dependencies

pkg/compose/build.go

@@ -139,6 +139,7 @@ func (s *composeService) ensureImagesExists(ctx context.Context, project *types.
 				project.Services[i].Labels = types.Labels{}
 			}
 			project.Services[i].CustomLabels[api.ImageDigestLabel] = digest
+			project.Services[i].CustomLabels[api.ImageNameLabel] = service.Image
 		}
 	}
 	return nil
@@ -191,6 +192,7 @@ func (s *composeService) getLocalImagesDigests(ctx context.Context, project *typ
 		digest, ok := images[imgName]
 		if ok {
 			project.Services[i].CustomLabels.Add(api.ImageDigestLabel, digest)
+			project.Services[i].CustomLabels.Add(api.ImageNameLabel, project.Services[i].Image)
 		}
 	}
 

pkg/compose/compose.go

@@ -130,9 +130,13 @@ func (s *composeService) projectFromName(containers Containers, projectName stri
 		serviceLabel := c.Labels[api.ServiceLabel]
 		_, ok := set[serviceLabel]
 		if !ok {
+			serviceImage := c.Image
+			if serviceNameFromLabel, ok := c.Labels[api.ImageNameLabel]; ok {
+				serviceImage = serviceNameFromLabel
+			}
 			set[serviceLabel] = &types.ServiceConfig{
 				Name:   serviceLabel,
-				Image:  c.Image,
+				Image:  serviceImage,
 				Labels: c.Labels,
 			}
 		}

pkg/compose/down.go

@@ -120,7 +120,7 @@ func (s *composeService) ensureVolumesDown(ctx context.Context, project *types.P
 
 func (s *composeService) ensureImagesDown(ctx context.Context, project *types.Project, options api.DownOptions, w progress.Writer) []downOp {
 	var ops []downOp
-	for image := range s.getServiceImages(options, project) {
+	for image := range s.getServiceImagesToRemove(options, project) {
 		image := image
 		ops = append(ops, func() error {
 			return s.removeImage(ctx, image, w)
@@ -190,16 +190,14 @@ func (s *composeService) removeNetwork(ctx context.Context, name string, w progr
 	return nil
 }
 
-func (s *composeService) getServiceImages(options api.DownOptions, project *types.Project) map[string]struct{} {
+func (s *composeService) getServiceImagesToRemove(options api.DownOptions, project *types.Project) map[string]struct{} {
 	images := map[string]struct{}{}
 	for _, service := range project.Services {
-		image := service.Image
-		if options.Images == "local" && image != "" {
+		image, ok := service.Labels[api.ImageNameLabel] // Information on the compose file at the creation of the container
+		if !ok || (options.Images == "local" && image != "") {
 			continue
 		}
-		if image == "" {
-			image = api.GetImageNameOrDefault(service, project.Name)
-		}
+		image = api.GetImageNameOrDefault(service, project.Name)
 		images[image] = struct{}{}
 	}
 	return images

pkg/compose/down_test.go

@@ -142,3 +142,90 @@ func TestDownRemoveVolumes(t *testing.T) {
 	err := tested.Down(context.Background(), strings.ToLower(testProject), compose.DownOptions{Volumes: true})
 	assert.NilError(t, err)
 }
+
+func TestDownRemoveImageLocal(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	api := mocks.NewMockAPIClient(mockCtrl)
+	cli := mocks.NewMockCli(mockCtrl)
+	tested.dockerCli = cli
+	cli.EXPECT().Client().Return(api).AnyTimes()
+
+	container := testContainer("service1", "123", false)
+	container.Labels[compose.ImageNameLabel] = ""
+
+	api.EXPECT().ContainerList(gomock.Any(), projectFilterListOpt(false)).Return(
+		[]moby.Container{container}, nil)
+
+	api.EXPECT().VolumeList(gomock.Any(), filters.NewArgs(projectFilter(strings.ToLower(testProject)))).
+		Return(volume.VolumeListOKBody{
+			Volumes: []*moby.Volume{{Name: "myProject_volume"}},
+		}, nil)
+	api.EXPECT().NetworkList(gomock.Any(), moby.NetworkListOptions{Filters: filters.NewArgs(projectFilter(strings.ToLower(testProject)))}).
+		Return(nil, nil)
+
+	api.EXPECT().ContainerStop(gomock.Any(), "123", nil).Return(nil)
+	api.EXPECT().ContainerRemove(gomock.Any(), "123", moby.ContainerRemoveOptions{Force: true}).Return(nil)
+
+	api.EXPECT().ImageRemove(gomock.Any(), "testproject-service1", moby.ImageRemoveOptions{}).Return(nil, nil)
+
+	err := tested.Down(context.Background(), strings.ToLower(testProject), compose.DownOptions{Images: "local"})
+	assert.NilError(t, err)
+}
+
+func TestDownRemoveImageLocalNoLabel(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	api := mocks.NewMockAPIClient(mockCtrl)
+	cli := mocks.NewMockCli(mockCtrl)
+	tested.dockerCli = cli
+	cli.EXPECT().Client().Return(api).AnyTimes()
+
+	container := testContainer("service1", "123", false)
+
+	api.EXPECT().ContainerList(gomock.Any(), projectFilterListOpt(false)).Return(
+		[]moby.Container{container}, nil)
+
+	api.EXPECT().VolumeList(gomock.Any(), filters.NewArgs(projectFilter(strings.ToLower(testProject)))).
+		Return(volume.VolumeListOKBody{
+			Volumes: []*moby.Volume{{Name: "myProject_volume"}},
+		}, nil)
+	api.EXPECT().NetworkList(gomock.Any(), moby.NetworkListOptions{Filters: filters.NewArgs(projectFilter(strings.ToLower(testProject)))}).
+		Return(nil, nil)
+
+	api.EXPECT().ContainerStop(gomock.Any(), "123", nil).Return(nil)
+	api.EXPECT().ContainerRemove(gomock.Any(), "123", moby.ContainerRemoveOptions{Force: true}).Return(nil)
+
+	err := tested.Down(context.Background(), strings.ToLower(testProject), compose.DownOptions{Images: "local"})
+	assert.NilError(t, err)
+}
+
+func TestDownRemoveImageAll(t *testing.T) {
+	mockCtrl := gomock.NewController(t)
+	defer mockCtrl.Finish()
+
+	api := mocks.NewMockAPIClient(mockCtrl)
+	cli := mocks.NewMockCli(mockCtrl)
+	tested.dockerCli = cli
+	cli.EXPECT().Client().Return(api).AnyTimes()
+
+	api.EXPECT().ContainerList(gomock.Any(), projectFilterListOpt(false)).Return(
+		[]moby.Container{testContainer("service1", "123", false)}, nil)
+
+	api.EXPECT().VolumeList(gomock.Any(), filters.NewArgs(projectFilter(strings.ToLower(testProject)))).
+		Return(volume.VolumeListOKBody{
+			Volumes: []*moby.Volume{{Name: "myProject_volume"}},
+		}, nil)
+	api.EXPECT().NetworkList(gomock.Any(), moby.NetworkListOptions{Filters: filters.NewArgs(projectFilter(strings.ToLower(testProject)))}).
+		Return(nil, nil)
+
+	api.EXPECT().ContainerStop(gomock.Any(), "123", nil).Return(nil)
+	api.EXPECT().ContainerRemove(gomock.Any(), "123", moby.ContainerRemoveOptions{Force: true}).Return(nil)
+
+	api.EXPECT().ImageRemove(gomock.Any(), "service1-img", moby.ImageRemoveOptions{}).Return(nil, nil)
+
+	err := tested.Down(context.Background(), strings.ToLower(testProject), compose.DownOptions{Images: "all"})
+	assert.NilError(t, err)
+}

pkg/compose/kill_test.go

@@ -109,6 +109,7 @@ func containerLabels(service string, oneOff bool) map[string]string {
 	composefile := filepath.Join(workingdir, "compose.yaml")
 	labels := map[string]string{
 		compose.ServiceLabel:     service,
+		compose.ImageNameLabel:   service + "-img",
 		compose.ConfigFilesLabel: composefile,
 		compose.WorkingDirLabel:  workingdir,
 		compose.ProjectLabel:     strings.ToLower(testProject)}