Bug 4693: Taking screenshots of all open windows in case of a test

failure / Improved cleanup of opened windows after a test case finished

java/com.sap.sailing.selenium.test/src/com/sap/sailing/selenium/core/WebDriverWindow.java

@@ -50,4 +50,8 @@ public class WebDriverWindow {
         if(!handles.contains(this.handle))
             throw new WebDriverException("Window closed or not initialized"); //$NON-NLS-1$
     }
+    
+    public WebDriver getWebDriver() {
+        return driver;
+    }
 }

java/com.sap.sailing.selenium.test/src/com/sap/sailing/selenium/core/WindowManager.java

@@ -1,6 +1,9 @@
 package com.sap.sailing.selenium.core;
 
+import java.util.HashSet;
+import java.util.Set;
 import java.util.function.BiConsumer;
+import java.util.function.Consumer;
 import java.util.function.Supplier;
 
 import org.openqa.selenium.Dimension;
@@ -13,6 +16,8 @@ import org.openqa.selenium.WebDriver;
  *   Riccardo Nimser (D049941)
  */
 public class WindowManager {
+    private final WebDriverWindow defaultWindow;
+    private final Set<WebDriverWindow> allWindows = new HashSet<>();
     private final WebDriver driver;
 
     private final Supplier<WebDriver> webDriverFactory;
@@ -25,28 +30,28 @@ public class WindowManager {
      */
     public WindowManager(WebDriver driver, Supplier<WebDriver> webDriverFactory) {
         this.driver = driver;
+        defaultWindow = new ManagedWebDriverWindow(this.driver, this.driver.getWindowHandle());
         this.webDriverFactory = webDriverFactory;
         
         setWindowMaximized(this.driver);
     }
     
     public void withExtraWindow(BiConsumer<WebDriverWindow, WebDriverWindow> defaultAndExtraWindow) {
-        final WebDriverWindow defaultWindow = new WebDriverWindow(this.driver, this.driver.getWindowHandle());
         final WebDriver extraDriver = webDriverFactory.get();
-        final WebDriverWindow extraWindow = new WebDriverWindow(extraDriver, extraDriver.getWindowHandle());
+        final WebDriverWindow extraWindow = new ManagedWebDriverWindow(extraDriver, extraDriver.getWindowHandle());
         
         extraWindow.switchToWindow();
         setWindowMaximized(extraDriver);
         defaultWindow.switchToWindow();
+        
+        defaultAndExtraWindow.accept(defaultWindow, extraWindow);
         try {
-            defaultAndExtraWindow.accept(defaultWindow, extraWindow);
-        } finally {
-            try {
-                extraDriver.quit();
-            } catch (Exception e) {
-                // This call may fail depending on the WebDriver being used
-            }
-            defaultWindow.switchToWindow();
+            // quit is explicitly not called in a finally block to ensure that both windows are still open
+            // when trying to create screenshots in case an error occurs
+            extraWindow.close();
+            extraDriver.quit();
+        } catch (Exception e) {
+            // This call may fail depending on the WebDriver being used
         }
     }
     
@@ -64,4 +69,28 @@ public class WindowManager {
             }
         }
     }
+    
+    public void forEachOpenedWindow(Consumer<WebDriverWindow> windowConsumer) {
+        this.allWindows.forEach(windowConsumer);
+    }
+    
+    public void closeAllExtraWindows() {
+        this.allWindows.forEach(window -> {
+            if (window != defaultWindow) {
+                window.close();
+            }
+        });
+    }
+    
+    private class ManagedWebDriverWindow extends WebDriverWindow {
+        protected ManagedWebDriverWindow(WebDriver driver, String handle) {
+            super(driver, handle);
+            allWindows.add(this);
+        }
+        @Override
+        public void close() {
+            allWindows.remove(this);
+            super.close();
+        }
+    }
 }

java/com.sap.sailing.selenium.test/src/com/sap/sailing/selenium/test/AbstractSeleniumTest.java

@@ -172,7 +172,7 @@ public abstract class AbstractSeleniumTest {
     //    }
     //}
     
-    private class ScreenShotRule extends TestWatchman {
+    private class ScreenShotAndCloseWindowRule extends TestWatchman {
         @Override
         public void failed(Throwable cause, FrameworkMethod method) {
             try {
@@ -181,13 +181,22 @@ public abstract class AbstractSeleniumTest {
                 exception.printStackTrace();
             }
         }
+        
+        @Override
+        public void finished(FrameworkMethod method) {
+            try {
+                environment.getWindowManager().closeAllExtraWindows();
+            } finally {
+                super.finished(method);
+            }
+        }
     }
     
     /**
      * <p>Rule for capturing of a screenshot if a test fails.</p>
      */
     @Rule
-    public final ScreenShotRule takeScreenshoot = new ScreenShotRule(/*generator*/);
+    public final ScreenShotAndCloseWindowRule takeScreenshotAndCloseWindows = new ScreenShotAndCloseWindowRule(/*generator*/);
 
     /**
      * <p>The test environment used for the execution of the the tests.</p>
@@ -243,25 +252,27 @@ public abstract class AbstractSeleniumTest {
     protected void captureScreenshot(String filename) {
         File screenshotFolder = this.environment.getScreenshotFolder();
         if (screenshotFolder != null) {
-            WebDriver driver = getWebDriver();
-            if (RemoteWebDriver.class.equals(driver.getClass())) {
-                driver = new Augmenter().augment(driver);
-            }
-            InputStream source = getScreenshotNotSupportedImage();
-            if (driver instanceof TakesScreenshot) {
-                source = new ByteArrayInputStream(((TakesScreenshot) driver).getScreenshotAs(OutputType.BYTES));
-            }
-            try {
-                File destinationDir = new File(screenshotFolder, getClass().getName());
-                destinationDir.mkdirs();
-                File destination = new File(destinationDir, filename + SCREENSHOT_FILE_EXTENSION); //$NON-NLS-1$
-                Path path = destination.toPath();
-                Files.copy(source, path, StandardCopyOption.REPLACE_EXISTING);
-                // ATTENTION: Do not remove this line because it is needed for the JUnit Attachment Plugin!
-                System.out.println(String.format(ATTACHMENT_FORMAT, destination.getCanonicalFile().toURI()));
-            } catch (IOException exception) {
-                throw new RuntimeException(exception);
-            }
+            this.environment.getWindowManager().forEachOpenedWindow(window -> {
+                WebDriver driver = window.getWebDriver();
+                if (RemoteWebDriver.class.equals(driver.getClass())) {
+                    driver = new Augmenter().augment(driver);
+                }
+                InputStream source = getScreenshotNotSupportedImage();
+                if (driver instanceof TakesScreenshot) {
+                    source = new ByteArrayInputStream(((TakesScreenshot) driver).getScreenshotAs(OutputType.BYTES));
+                }
+                try {
+                    File destinationDir = new File(screenshotFolder, getClass().getName());
+                    destinationDir.mkdirs();
+                    File destination = new File(destinationDir, filename + SCREENSHOT_FILE_EXTENSION); //$NON-NLS-1$
+                    Path path = destination.toPath();
+                    Files.copy(source, path, StandardCopyOption.REPLACE_EXISTING);
+                    // ATTENTION: Do not remove this line because it is needed for the JUnit Attachment Plugin!
+                    System.out.println(String.format(ATTACHMENT_FORMAT, destination.getCanonicalFile().toURI()));
+                } catch (IOException exception) {
+                    throw new RuntimeException(exception);
+                }
+            });
         }
     }