diff --git a/umap/tests/integration/test_websocket_sync.py b/umap/tests/integration/test_websocket_sync.py
index 604bdc6a..a038f784 100644
--- a/umap/tests/integration/test_websocket_sync.py
+++ b/umap/tests/integration/test_websocket_sync.py
@@ -69,8 +69,8 @@ def test_websocket_connection_can_sync_markers(
 
     # Delete a marker from peer A and check it's been deleted on peer B
     a_first_marker.click(button="right")
-    peerA.on("dialog", lambda dialog: dialog.accept())
     peerA.get_by_role("link", name="Delete this feature").click()
+    peerA.locator("dialog").get_by_role("button", name="OK").click()
     expect(a_marker_pane).to_have_count(1)
     expect(b_marker_pane).to_have_count(1)
 
@@ -153,8 +153,8 @@ def test_websocket_connection_can_sync_polygons(
 
     # Delete a polygon from peer A and check it's been deleted on peer B
     a_polygon.click(button="right")
-    peerA.on("dialog", lambda dialog: dialog.accept())
     peerA.get_by_role("link", name="Delete this feature").click()
+    peerA.locator("dialog").get_by_role("button", name="OK").click()
     expect(a_polygons).to_have_count(0)
     expect(b_polygons).to_have_count(0)