Reimplement cleanup logic

Co-Authored-By: Johannes Rabauer <[email protected]>

Author: AB-xdev

testcontainers-advanced-imagebuilder/src/main/java/software/xdev/testcontainers/imagebuilder/AdvancedImageFromDockerFile.java

@@ -149,12 +149,6 @@ protected String resolve()
 
 		try
 		{
-			if(this.deleteOnExit)
-			{
-				this.log().info("Registering image for cleanup when tests are finished");
-				ResourceReaper.instance().registerImageForCleanup(this.dockerImageName);
-			}
-			
 			// We have to use pipes to avoid high memory consumption since users might want to build huge images
 			final PipedInputStream in = new PipedInputStream();
 			final PipedOutputStream out = new PipedOutputStream(in);
@@ -168,6 +162,7 @@ protected String resolve()
 			{
 				labels.putAll(buildImageCmd.getLabels());
 			}
+			this.deleteImageOnExitIfRequired(labels);
 			labels.putAll(DockerClientFactory.DEFAULT_LABELS);
 			buildImageCmd.withLabels(labels);
 
@@ -209,6 +204,18 @@ protected String resolve()
 		}
 	}
 
+	@SuppressWarnings("deprecation") // There is no alternative and it's also used in the default implementation
+	protected void deleteImageOnExitIfRequired(final Map<String, String> labels)
+	{
+		if(!this.deleteOnExit)
+		{
+			return;
+		}
+		
+		this.log().debug("Registering image for cleanup when finished");
+		labels.putAll(ResourceReaper.instance().getLabels());
+	}
+	
 	protected long getBytesToDockerDaemon(final PipedOutputStream out) throws IOException
 	{
 		long bytesToDockerDaemon = 0L;