From 2de6e0383c2d14b39e7ac900f220ff3665db0dbe Mon Sep 17 00:00:00 2001
From: Mark Paluch <mpaluch@vmware.com>
Date: Fri, 30 Apr 2021 10:45:05 +0200
Subject: [PATCH] Polishing.

Use Awaitility to ensure Cassandra connectivity. Start Cassandra container only once.

See #606.
---
 .../OptimisticPersonRepositoryTests.java      | 13 ++++++----
 cassandra/util/pom.xml                        |  5 ++++
 .../cassandra/util/CassandraExtension.java    | 26 ++++++++++++++++---
 geode/pom.xml                                 |  5 ----
 pom.xml                                       |  6 +++++
 5 files changed, 42 insertions(+), 13 deletions(-)

diff --git a/cassandra/example/src/test/java/example/springdata/cassandra/optimisticlocking/OptimisticPersonRepositoryTests.java b/cassandra/example/src/test/java/example/springdata/cassandra/optimisticlocking/OptimisticPersonRepositoryTests.java
index a8693beb..4266f036 100644
--- a/cassandra/example/src/test/java/example/springdata/cassandra/optimisticlocking/OptimisticPersonRepositoryTests.java
+++ b/cassandra/example/src/test/java/example/springdata/cassandra/optimisticlocking/OptimisticPersonRepositoryTests.java
@@ -17,6 +17,8 @@ package example.springdata.cassandra.optimisticlocking;
 
 import static org.assertj.core.api.Assertions.*;
 
+import example.springdata.cassandra.util.CassandraKeyspace;
+
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
@@ -33,13 +35,14 @@ import org.springframework.data.cassandra.core.query.Criteria;
  * @author Mark Paluch
  */
 @SpringBootTest(classes = BasicConfiguration.class)
-public class OptimisticPersonRepositoryTests {
+@CassandraKeyspace
+class OptimisticPersonRepositoryTests {
 
 	@Autowired OptimisticPersonRepository repository;
 	@Autowired CassandraOperations operations;
 
 	@BeforeEach
-	public void setUp() {
+	void setUp() {
 		repository.deleteAll();
 	}
 
@@ -48,7 +51,7 @@ public class OptimisticPersonRepositoryTests {
 	 * and increment the version property.
 	 */
 	@Test
-	public void insertShouldIncrementVersion() {
+	void insertShouldIncrementVersion() {
 
 		var person = new OptimisticPerson(42L, 0, "Walter White");
 
@@ -61,7 +64,7 @@ public class OptimisticPersonRepositoryTests {
 	 * Modifying an existing object will update the last modified fields.
 	 */
 	@Test
-	public void updateShouldDetectChangedEntity() {
+	void updateShouldDetectChangedEntity() {
 
 		var person = new OptimisticPerson(42L, 0, "Walter White");
 
@@ -85,7 +88,7 @@ public class OptimisticPersonRepositoryTests {
 	 * statement through {@link CassandraOperations#update(Object, UpdateOptions)}.
 	 */
 	@Test
-	public void updateUsingLightWeightTransactions() {
+	void updateUsingLightWeightTransactions() {
 
 		var person = new SimplePerson();
 		person.setId(42L);
diff --git a/cassandra/util/pom.xml b/cassandra/util/pom.xml
index fc88b10a..adf66b17 100644
--- a/cassandra/util/pom.xml
+++ b/cassandra/util/pom.xml
@@ -35,6 +35,11 @@
 			<artifactId>java-driver-core</artifactId>
 		</dependency>
 
+		<dependency>
+			<groupId>org.awaitility</groupId>
+			<artifactId>awaitility</artifactId>
+		</dependency>
+
 	</dependencies>
 
 </project>
diff --git a/cassandra/util/src/main/java/example/springdata/cassandra/util/CassandraExtension.java b/cassandra/util/src/main/java/example/springdata/cassandra/util/CassandraExtension.java
index 3b1a54ff..1023fc04 100644
--- a/cassandra/util/src/main/java/example/springdata/cassandra/util/CassandraExtension.java
+++ b/cassandra/util/src/main/java/example/springdata/cassandra/util/CassandraExtension.java
@@ -17,7 +17,9 @@ package example.springdata.cassandra.util;
 
 import java.net.InetSocketAddress;
 import java.util.Optional;
+import java.util.concurrent.Callable;
 
+import org.awaitility.Awaitility;
 import org.junit.jupiter.api.extension.BeforeAllCallback;
 import org.junit.jupiter.api.extension.ExtensionContext;
 import org.junit.platform.commons.util.AnnotationUtils;
@@ -39,6 +41,8 @@ class CassandraExtension implements BeforeAllCallback {
 	private static final ExtensionContext.Namespace NAMESPACE = ExtensionContext.Namespace
 			.create(CassandraExtension.class);
 
+	private static CassandraContainer container;
+
 	@Override
 	public void beforeAll(ExtensionContext context) {
 
@@ -57,10 +61,22 @@ class CassandraExtension implements BeforeAllCallback {
 
 		keyspace.before();
 
+
+		Callable<CqlSession> sessionFactory = () -> CqlSession.builder()
+				.addContactPoint(new InetSocketAddress(keyspace.host(), keyspace.port())).withLocalDatacenter("datacenter1")
+				.build();
+		Awaitility.await().ignoreExceptions().untilAsserted(() -> {
+
+			sessionFactory.call().close();
+		});
+
 		var session = store.getOrComputeIfAbsent(CqlSession.class, it -> {
 
-			return CqlSession.builder().addContactPoint(new InetSocketAddress(keyspace.host(), keyspace.port()))
-					.withLocalDatacenter("datacenter1").build();
+			try {
+				return sessionFactory.call();
+			} catch (Exception e) {
+				throw new IllegalStateException(e);
+			}
 		}, CqlSession.class);
 
 		session.execute(String.format("CREATE KEYSPACE IF NOT EXISTS %s \n"
@@ -78,7 +94,11 @@ class CassandraExtension implements BeforeAllCallback {
 
 	private CassandraContainer<?> runTestcontainer() {
 
-		var container = new CassandraContainer<>(getCassandraDockerImageName());
+		if (container != null) {
+			return container;
+		}
+
+		container = new CassandraContainer<>(getCassandraDockerImageName());
 		container.withReuse(true);
 
 		container.start();
diff --git a/geode/pom.xml b/geode/pom.xml
index fe6bdea8..72e20edb 100755
--- a/geode/pom.xml
+++ b/geode/pom.xml
@@ -36,11 +36,6 @@
 
 	<dependencyManagement>
 		<dependencies>
-			<dependency>
-				<groupId>org.awaitility</groupId>
-				<artifactId>awaitility</artifactId>
-				<version>4.0.2</version>
-			</dependency>
 			<dependency>
 				<groupId>org.iq80.snappy</groupId>
 				<artifactId>snappy</artifactId>
diff --git a/pom.xml b/pom.xml
index cf7430c3..373810c6 100644
--- a/pom.xml
+++ b/pom.xml
@@ -111,6 +111,12 @@
 				<artifactId>spring-data-keyvalue</artifactId>
 				<version>2.5.1-SNAPSHOT</version>
 			</dependency>
+
+			<dependency>
+				<groupId>org.awaitility</groupId>
+				<artifactId>awaitility</artifactId>
+				<version>4.0.2</version>
+			</dependency>
 		</dependencies>
 	</dependencyManagement>