diff --git a/test/unit/org/openstreetmap/josm/plugins/mapwithai/backend/MapWithAIMoveActionTest.java b/test/unit/org/openstreetmap/josm/plugins/mapwithai/backend/MapWithAIMoveActionTest.java
index 5211a6d..92fbcd2 100644
--- a/test/unit/org/openstreetmap/josm/plugins/mapwithai/backend/MapWithAIMoveActionTest.java
+++ b/test/unit/org/openstreetmap/josm/plugins/mapwithai/backend/MapWithAIMoveActionTest.java
@@ -7,6 +7,10 @@ import static org.junit.jupiter.api.Assertions.assertNotNull;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
+import java.util.Arrays;
+import java.util.HashMap;
+import java.util.Map;
+
 import javax.swing.JOptionPane;
 
 import org.awaitility.Awaitility;
@@ -30,6 +34,7 @@ import org.openstreetmap.josm.plugins.mapwithai.backend.commands.conflation.Dupl
 import org.openstreetmap.josm.testutils.JOSMTestRules;
 import org.openstreetmap.josm.testutils.mockers.JOptionPaneSimpleMocker;
 import org.openstreetmap.josm.testutils.mockers.WindowMocker;
+import org.openstreetmap.josm.tools.I18n;
 
 import com.google.common.collect.ImmutableMap;
 
@@ -145,6 +150,14 @@ public class MapWithAIMoveActionTest {
     public void testMaxAddNotification() {
         TestUtils.assumeWorkingJMockit();
         new WindowMocker();
+        Map<String, Object> map = new HashMap<>();
+        for (String text : Arrays.asList(I18n.marktr("Sequence: Merge {0} nodes"), I18n.marktr("Delete {0} nodes"))) {
+            for (int i = 0; i < 10; i++) {
+                map.computeIfAbsent(I18n.tr(text, i), (t) -> JOptionPane.NO_OPTION);
+            }
+        }
+        new JOptionPaneSimpleMocker(map);
+
         NotificationMocker notification = new NotificationMocker();
         DataSet ds = MapWithAIDataUtils.getLayer(true).getDataSet();
         MainApplication.getLayerManager().addLayer(new OsmDataLayer(new DataSet(), "TEST", null));
@@ -153,7 +166,8 @@ public class MapWithAIMoveActionTest {
             ds.addPrimitive(new Node(LatLon.ZERO));
         }
         for (int i = 0; i < 11; i++) {
-            GuiHelper.runInEDTAndWait(() -> ds.setSelected(ds.allNonDeletedPrimitives().iterator().next()));
+            GuiHelper
+                    .runInEDTAndWaitWithException(() -> ds.setSelected(ds.allNonDeletedPrimitives().iterator().next()));
             moveAction.actionPerformed(null);
         }
         assertTrue(notification.shown);
diff --git a/test/unit/org/openstreetmap/josm/plugins/mapwithai/commands/MapWithAIAddComandTest.java b/test/unit/org/openstreetmap/josm/plugins/mapwithai/commands/MapWithAIAddComandTest.java
index ed5cf09..78e20cc 100644
--- a/test/unit/org/openstreetmap/josm/plugins/mapwithai/commands/MapWithAIAddComandTest.java
+++ b/test/unit/org/openstreetmap/josm/plugins/mapwithai/commands/MapWithAIAddComandTest.java
@@ -10,11 +10,11 @@ import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import java.io.File;
 import java.lang.reflect.InvocationTargetException;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.Future;
-import java.util.concurrent.FutureTask;
 
 import javax.swing.JOptionPane;
 
@@ -40,9 +40,11 @@ import org.openstreetmap.josm.gui.layer.OsmDataLayer;
 import org.openstreetmap.josm.gui.progress.NullProgressMonitor;
 import org.openstreetmap.josm.plugins.mapwithai.backend.commands.conflation.ConnectedCommand;
 import org.openstreetmap.josm.plugins.mapwithai.testutils.MapWithAITestRules;
+import org.openstreetmap.josm.plugins.mapwithai.testutils.PleaseWaitDialogMocker;
 import org.openstreetmap.josm.plugins.mapwithai.testutils.SwingUtilitiesMocker;
 import org.openstreetmap.josm.testutils.JOSMTestRules;
 import org.openstreetmap.josm.testutils.mockers.JOptionPaneSimpleMocker;
+import org.openstreetmap.josm.testutils.mockers.WindowMocker;
 import org.openstreetmap.josm.tools.Logging;
 
 import com.google.common.collect.ImmutableMap;
@@ -103,6 +105,7 @@ public class MapWithAIAddComandTest {
 
     @Test
     public void testCreateConnections() {
+        new PleaseWaitDialogMocker();
         final DataSet ds1 = new DataSet();
         final Way way1 = TestUtils.newWay(HIGHWAY_RESIDENTIAL, new Node(new LatLon(0, 0)),
                 new Node(new LatLon(0, 0.15)));
@@ -217,7 +220,12 @@ public class MapWithAIAddComandTest {
     @Test
     public void testRegression18351() throws InvocationTargetException, InterruptedException {
         System.getProperty("java.awt.headless", "true");
-        List<FutureTask<?>> futures = new SwingUtilitiesMocker().futureTasks;
+        TestUtils.assumeWorkingJMockit();
+        new WindowMocker();
+        new PleaseWaitDialogMocker();
+        SwingUtilitiesMocker swingMocker = new SwingUtilitiesMocker();
+
+        swingMocker.futureTasks.clear();
         Way way = TestUtils.newWay("highway=residential mapwithai:source=MapWithAI source=digitalglobe",
                 new Node(new LatLon(39.0339521, -108.4874581)), new Node(new LatLon(39.0292629, -108.4875117)));
         way.firstNode().put("dupe", "n176220609");
@@ -239,12 +247,14 @@ public class MapWithAIAddComandTest {
         assertNotNull(osmData.getPrimitiveById(176220609, OsmPrimitiveType.NODE));
         assertNotNull(osmData.getPrimitiveById(way));
         Awaitility.await().pollDelay(Durations.ONE_HUNDRED_MILLISECONDS);
+        List<Future<?>> futures = new ArrayList<>(swingMocker.futureTasks);
         Assertions.assertDoesNotThrow(() -> {
             for (Future<?> future : futures) {
                 try {
                     future.get();
                 } catch (Exception e) {
                     Logging.error(e);
+                    Logging.logWithStackTrace(Logging.LEVEL_ERROR, e);
                     throw e;
                 }
             }
diff --git a/test/unit/org/openstreetmap/josm/plugins/mapwithai/testutils/PleaseWaitDialogMocker.java b/test/unit/org/openstreetmap/josm/plugins/mapwithai/testutils/PleaseWaitDialogMocker.java
new file mode 100644
index 0000000..ee0889a
--- /dev/null
+++ b/test/unit/org/openstreetmap/josm/plugins/mapwithai/testutils/PleaseWaitDialogMocker.java
@@ -0,0 +1,31 @@
+// License: GPL. For details, see LICENSE file.
+package org.openstreetmap.josm.plugins.mapwithai.testutils;
+
+import java.awt.Component;
+
+import org.openstreetmap.josm.gui.PleaseWaitDialog;
+
+import mockit.Mock;
+import mockit.MockUp;
+
+public class PleaseWaitDialogMocker extends MockUp<PleaseWaitDialog> {
+    @Mock
+    protected void adjustLayout() {
+        // We don't want to adjust the layout...
+    }
+
+    @Mock
+    public void setLocationRelativeTo(Component c) {
+        // Do nothing
+    }
+
+    @Mock
+    public void setVisible(boolean visible) {
+        // Do nothing
+    }
+
+    @Mock
+    public void dispose() {
+        // Do nothing...
+    }
+}