fix cleanup for rolling updates (#706)

internal/actions/update.go

@@ -73,6 +73,7 @@ func performRollingRestart(containers []container.Container, client container.Cl
 		if containers[i].Stale {
 			stopStaleContainer(containers[i], client, params)
 			restartStaleContainer(containers[i], client, params)
+			cleanupImageIDs[containers[i].ImageID()] = true
 		}
 	}
 

internal/actions/update_test.go

@@ -80,6 +80,14 @@ var _ = Describe("the update action", func() {
 				Expect(client.TestData.TriedToRemoveImageCount).To(Equal(2))
 			})
 		})
+		When("performing a rolling restart update", func() {
+			It("should try to remove the image once", func() {
+
+				err := actions.Update(client, types.UpdateParams{Cleanup: true, RollingRestart: true})
+				Expect(err).NotTo(HaveOccurred())
+				Expect(client.TestData.TriedToRemoveImageCount).To(Equal(1))
+			})
+		})
 	})
 
 	When("watchtower has been instructed to monitor only", func() {