--- a/test/jdk/ProblemList.txt	Tue Nov 14 19:33:37 2017 -0800
+++ b/test/jdk/ProblemList.txt	Wed Nov 15 09:31:17 2017 -0800
@@ -245,8 +245,6 @@
 
 # jdk_swing
 
-sanity/client/SwingSet/src/ButtonDemoScreenshotTest.java        8157338 generic-all
-
 ############################################################################
 
 # jdk_text

--- a/test/jdk/sanity/client/SwingSet/src/ButtonDemoScreenshotTest.java	Tue Nov 14 19:33:37 2017 -0800
+++ b/test/jdk/sanity/client/SwingSet/src/ButtonDemoScreenshotTest.java	Wed Nov 15 09:31:17 2017 -0800
@@ -23,18 +23,21 @@
 
 import com.sun.swingset3.demos.button.ButtonDemo;
 import org.jtregext.GuiTestListener;
-import java.awt.Point;
-import java.awt.Robot;
-import java.awt.event.InputEvent;
-import java.awt.image.BufferedImage;
 import org.netbeans.jemmy.ClassReference;
+import org.netbeans.jemmy.ComponentChooser;
 import org.netbeans.jemmy.image.StrictImageComparator;
 import org.netbeans.jemmy.operators.JButtonOperator;
 import org.netbeans.jemmy.operators.JFrameOperator;
-import static org.jemmy2ext.JemmyExt.*;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Listeners;
 import org.testng.annotations.Test;
-import static com.sun.swingset3.demos.button.ButtonDemo.*;
-import org.testng.annotations.Listeners;
+
+import java.awt.Component;
+import java.awt.Robot;
+import java.awt.image.BufferedImage;
+
+import static com.sun.swingset3.demos.button.ButtonDemo.DEMO_TITLE;
+import static org.jemmy2ext.JemmyExt.*;
 
 /*
  * @test
@@ -55,7 +58,13 @@
 @Listeners(GuiTestListener.class)
 public class ButtonDemoScreenshotTest {
 
-    private static final int BUTTON_COUNT = 6; // TODO: Decide about "open browser" buttons (value was 8 originally)
+    private static final int[] BUTTONS = {0, 1, 2, 3, 4, 5}; // "open browser" buttons (6, 7) open a browser, so ignore
+    private static StrictImageComparator sComparator = null;
+
+    @BeforeClass
+    public void init() {
+        sComparator = new StrictImageComparator();
+    }
 
     @Test
     public void test() throws Exception {
@@ -67,32 +76,37 @@
         waitImageIsStill(rob, mainFrame);
 
         // Check all the buttons
-        for (int i = 0; i < BUTTON_COUNT; i++) {
+        for (int i : BUTTONS) {
             checkButton(mainFrame, i, rob);
         }
     }
 
-    public void checkButton(JFrameOperator jfo, int i, Robot rob) {
+    private void checkButton(JFrameOperator jfo, int i, Robot rob) {
         JButtonOperator button = new JButtonOperator(jfo, i);
-
-        Point loc = button.getLocationOnScreen();
-        rob.mouseMove(loc.x, loc.y);
+        button.moveMouse(button.getCenterX(), button.getCenterY());
 
         BufferedImage initialButtonImage = capture(rob, button);
         assertNotBlack(initialButtonImage);
-        save(initialButtonImage, "button" + i + "_0initial.png");
-        rob.mousePress(InputEvent.BUTTON1_MASK);
+        save(initialButtonImage, "button" + i + ".png");
+
+        BufferedImage[] pressedImage = new BufferedImage[1];
+
+        button.pressMouse();
         try {
             waitPressed(button);
-            BufferedImage pressedButtonImage = capture(rob, button);
-            assertNotBlack(pressedButtonImage);
-            save(pressedButtonImage, "button" + i + "_1pressed.png");
-
-            StrictImageComparator sComparator = new StrictImageComparator();
-            assertNotEquals("Button " + i + " Test", sComparator, initialButtonImage, pressedButtonImage);
+            button.waitState(new ComponentChooser() {
+                public boolean checkComponent(Component c) {
+                    pressedImage[0] = capture(rob, button);
+                    assertNotBlack(pressedImage[0]);
+                    return !sComparator.compare(initialButtonImage, pressedImage[0]);
+                }
+                public String getDescription() {
+                    return "Button with new image";
+                }
+            });
         } finally {
-            rob.mouseRelease(InputEvent.BUTTON1_MASK);
+            if(pressedImage[0] != null) save(pressedImage[0], "button" + i + "_pressed.png");
+            button.releaseMouse();
         }
     }
-
 }