From fc3b43e83257a49b64c278ea995ffc741c4c2a88 Mon Sep 17 00:00:00 2001
From: Madhura Bhave <mbhave@pivotal.io>
Date: Thu, 4 Jan 2018 14:26:34 -0800
Subject: [PATCH] Wait for neo4j to start accepting connections

See gh-10516
---
 .../neo4j/DataNeo4jTestIntegrationTests.java  | 33 ++++++++++++++++---
 ...TestWithIncludeFilterIntegrationTests.java |  2 +-
 2 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/spring-boot-project/spring-boot-test-autoconfigure/src/test/java/org/springframework/boot/test/autoconfigure/data/neo4j/DataNeo4jTestIntegrationTests.java b/spring-boot-project/spring-boot-test-autoconfigure/src/test/java/org/springframework/boot/test/autoconfigure/data/neo4j/DataNeo4jTestIntegrationTests.java
index 71e34b6748..5ded2a84e8 100644
--- a/spring-boot-project/spring-boot-test-autoconfigure/src/test/java/org/springframework/boot/test/autoconfigure/data/neo4j/DataNeo4jTestIntegrationTests.java
+++ b/spring-boot-project/spring-boot-test-autoconfigure/src/test/java/org/springframework/boot/test/autoconfigure/data/neo4j/DataNeo4jTestIntegrationTests.java
@@ -16,6 +16,8 @@
 
 package org.springframework.boot.test.autoconfigure.data.neo4j;
 
+import java.util.concurrent.Callable;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Supplier;
 
 import org.junit.ClassRule;
@@ -23,7 +25,11 @@ import org.junit.Rule;
 import org.junit.Test;
 import org.junit.rules.ExpectedException;
 import org.junit.runner.RunWith;
+import org.neo4j.ogm.config.Configuration;
 import org.neo4j.ogm.session.Session;
+import org.neo4j.ogm.session.SessionFactory;
+import org.rnorth.ducttape.TimeoutException;
+import org.rnorth.ducttape.unreliables.Unreliables;
 import org.testcontainers.containers.FixedHostPortGenericContainer;
 import org.testcontainers.containers.GenericContainer;
 import org.testcontainers.containers.wait.HostPortWaitStrategy;
@@ -49,7 +55,7 @@ public class DataNeo4jTestIntegrationTests {
 	@ClassRule
 	public static DockerTestContainer<GenericContainer> genericContainer = new DockerTestContainer<>((Supplier<GenericContainer>) () -> new FixedHostPortGenericContainer("neo4j:latest")
 			.withFixedExposedPort(7687, 7687)
-			.waitingFor(new AdditionalSleepWaitStrategy()).withEnv("NEO4J_AUTH", "none"));
+			.waitingFor(new ConnectionVerifyingWaitStrategy()).withEnv("NEO4J_AUTH", "none"));
 
 
 	@Rule
@@ -80,18 +86,35 @@ public class DataNeo4jTestIntegrationTests {
 		this.applicationContext.getBean(ExampleService.class);
 	}
 
-	static class AdditionalSleepWaitStrategy extends HostPortWaitStrategy {
+	static class ConnectionVerifyingWaitStrategy extends HostPortWaitStrategy {
 
 		@Override
 		protected void waitUntilReady() {
 			super.waitUntilReady();
+			Configuration configuration = new Configuration.Builder()
+					.uri("bolt://localhost:7687").build();
+			SessionFactory sessionFactory = new SessionFactory(configuration,
+					"org.springframework.boot.test.autoconfigure.data.neo4j");
 			try {
-				Thread.sleep(5000);
+				Unreliables.retryUntilTrue((int) startupTimeout.getSeconds(), TimeUnit.SECONDS,
+						checkConnection(sessionFactory));
 			}
-			catch (InterruptedException e) {
-				Thread.currentThread().interrupt();
+			catch (TimeoutException e) {
+				throw new IllegalStateException();
 			}
 		}
+
+		private Callable<Boolean> checkConnection(SessionFactory sessionFactory) {
+			return () -> {
+				try {
+					sessionFactory.openSession().beginTransaction().close();
+					return true;
+				}
+				catch (Exception ex) {
+					return false;
+				}
+			};
+		}
 	}
 
 }
diff --git a/spring-boot-project/spring-boot-test-autoconfigure/src/test/java/org/springframework/boot/test/autoconfigure/data/neo4j/DataNeo4jTestWithIncludeFilterIntegrationTests.java b/spring-boot-project/spring-boot-test-autoconfigure/src/test/java/org/springframework/boot/test/autoconfigure/data/neo4j/DataNeo4jTestWithIncludeFilterIntegrationTests.java
index b591f66d8b..0e18c46b1c 100644
--- a/spring-boot-project/spring-boot-test-autoconfigure/src/test/java/org/springframework/boot/test/autoconfigure/data/neo4j/DataNeo4jTestWithIncludeFilterIntegrationTests.java
+++ b/spring-boot-project/spring-boot-test-autoconfigure/src/test/java/org/springframework/boot/test/autoconfigure/data/neo4j/DataNeo4jTestWithIncludeFilterIntegrationTests.java
@@ -44,7 +44,7 @@ public class DataNeo4jTestWithIncludeFilterIntegrationTests {
 	@ClassRule
 	public static DockerTestContainer<GenericContainer> genericContainer = new DockerTestContainer<>((Supplier<GenericContainer>) () -> new FixedHostPortGenericContainer("neo4j:latest")
 			.withFixedExposedPort(7687, 7687)
-			.waitingFor(new DataNeo4jTestIntegrationTests.AdditionalSleepWaitStrategy()).withEnv("NEO4J_AUTH", "none"));
+			.waitingFor(new DataNeo4jTestIntegrationTests.ConnectionVerifyingWaitStrategy()).withEnv("NEO4J_AUTH", "none"));
 
 	@Autowired
 	private ExampleService service;