From e3922fd86aea38646d308c937510085cdb7ab6bf Mon Sep 17 00:00:00 2001
From: Love Kristofer Leifland <love.leifland@neotechnology.com>
Date: Thu, 28 Dec 2023 13:57:03 +0100
Subject: [PATCH] [tyXNOmR9] Investigate flaky test

---
 .../test/java/apoc/spatial/GeocodeTest.java   | 45 +++++++++----------
 1 file changed, 20 insertions(+), 25 deletions(-)

diff --git a/core/src/test/java/apoc/spatial/GeocodeTest.java b/core/src/test/java/apoc/spatial/GeocodeTest.java
index 79a8e8bd6..cb6e49811 100755
--- a/core/src/test/java/apoc/spatial/GeocodeTest.java
+++ b/core/src/test/java/apoc/spatial/GeocodeTest.java
@@ -20,17 +20,21 @@
 
 import static apoc.ApocConfig.apocConfig;
 import static apoc.util.MapUtil.map;
-import static org.junit.Assert.*;
-import static org.neo4j.test.assertion.Assert.assertEventually;
+import static org.awaitility.Awaitility.await;
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertNotNull;
+import static org.junit.Assert.assertThrows;
+import static org.junit.Assert.assertTrue;
 
 import apoc.util.JsonUtil;
 import apoc.util.TestUtil;
 import inet.ipaddr.IPAddressString;
 import java.io.InputStream;
+import java.time.Duration;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
-import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Consumer;
 import java.util.stream.Stream;
@@ -295,27 +299,18 @@ private void testGeocodeAddress(
 
     private void waitForServerResponseOK(
             String query, Map<String, Object> params, AtomicLong time, Consumer<Result> resultObjectFunction) {
-        assertEventually(
-                () -> {
-                    try {
-                        long start = System.currentTimeMillis();
-                        db.executeTransactionally(query, params, res -> {
-                            resultObjectFunction.accept(res);
-                            return null;
-                        });
-
-                        time.addAndGet(System.currentTimeMillis() - start);
-                        return true;
-                    } catch (Exception e) {
-                        String msg = e.getMessage();
-                        if (msg.contains("Server returned HTTP response code") || msg.contains("connect timed out")) {
-                            return false;
-                        }
-                        throw e;
-                    }
-                },
-                (value) -> value,
-                20L,
-                TimeUnit.SECONDS);
+        await("waitForServerResponseOK")
+                .atMost(Duration.ofSeconds(30))
+                .pollInterval(Duration.ofMillis(250))
+                .pollInSameThread()
+                .untilAsserted(() -> {
+                    long start = System.currentTimeMillis();
+                    db.executeTransactionally(query, params, res -> {
+                        resultObjectFunction.accept(res);
+                        return null;
+                    });
+
+                    time.addAndGet(System.currentTimeMillis() - start);
+                });
     }
 }