Improve Cassandra health indicator with more robust mechanism

See gh-23041

spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/cassandra/CassandraDriverHealthIndicator.java

@@ -16,14 +16,17 @@
 
 package org.springframework.boot.actuate.cassandra;
 
-import com.datastax.oss.driver.api.core.ConsistencyLevel;
+import java.util.Collection;
+import java.util.Optional;
+
 import com.datastax.oss.driver.api.core.CqlSession;
-import com.datastax.oss.driver.api.core.cql.Row;
-import com.datastax.oss.driver.api.core.cql.SimpleStatement;
+import com.datastax.oss.driver.api.core.metadata.Node;
+import com.datastax.oss.driver.api.core.metadata.NodeState;
 
 import org.springframework.boot.actuate.health.AbstractHealthIndicator;
 import org.springframework.boot.actuate.health.Health;
 import org.springframework.boot.actuate.health.HealthIndicator;
+import org.springframework.boot.actuate.health.Status;
 import org.springframework.util.Assert;
 
 /**
@@ -31,13 +34,11 @@ import org.springframework.util.Assert;
  * Cassandra data stores.
  *
  * @author Alexandre Dutra
+ * @author Tomasz Lelek
  * @since 2.4.0
  */
 public class CassandraDriverHealthIndicator extends AbstractHealthIndicator {
 
-	private static final SimpleStatement SELECT = SimpleStatement
-			.newInstance("SELECT release_version FROM system.local").setConsistencyLevel(ConsistencyLevel.LOCAL_ONE);
-
 	private final CqlSession session;
 
 	/**
@@ -52,11 +53,10 @@ public class CassandraDriverHealthIndicator extends AbstractHealthIndicator {
 
 	@Override
 	protected void doHealthCheck(Health.Builder builder) throws Exception {
-		Row row = this.session.execute(SELECT).one();
-		builder.up();
-		if (row != null && !row.isNull(0)) {
-			builder.withDetail("version", row.getString(0));
-		}
+		Collection<Node> nodes = this.session.getMetadata().getNodes().values();
+		Optional<Node> nodeUp = nodes.stream().filter((node) -> node.getState() == NodeState.UP).findAny();
+		builder.status(nodeUp.isPresent() ? Status.UP : Status.DOWN);
+		nodeUp.map(Node::getCassandraVersion).ifPresent((version) -> builder.withDetail("version", version));
 	}
 
 }

spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/cassandra/CassandraDriverReactiveHealthIndicator.java

@@ -15,14 +15,18 @@
  */
 package org.springframework.boot.actuate.cassandra;
 
-import com.datastax.oss.driver.api.core.ConsistencyLevel;
+import java.util.Collection;
+import java.util.Optional;
+
 import com.datastax.oss.driver.api.core.CqlSession;
-import com.datastax.oss.driver.api.core.cql.SimpleStatement;
+import com.datastax.oss.driver.api.core.metadata.Node;
+import com.datastax.oss.driver.api.core.metadata.NodeState;
 import reactor.core.publisher.Mono;
 
 import org.springframework.boot.actuate.health.AbstractReactiveHealthIndicator;
 import org.springframework.boot.actuate.health.Health;
 import org.springframework.boot.actuate.health.ReactiveHealthIndicator;
+import org.springframework.boot.actuate.health.Status;
 import org.springframework.util.Assert;
 
 /**
@@ -30,13 +34,11 @@ import org.springframework.util.Assert;
  * for Cassandra data stores.
  *
  * @author Alexandre Dutra
+ * @author Tomasz Lelek
  * @since 2.4.0
  */
 public class CassandraDriverReactiveHealthIndicator extends AbstractReactiveHealthIndicator {
 
-	private static final SimpleStatement SELECT = SimpleStatement
-			.newInstance("SELECT release_version FROM system.local").setConsistencyLevel(ConsistencyLevel.LOCAL_ONE);
-
 	private final CqlSession session;
 
 	/**
@@ -51,8 +53,13 @@ public class CassandraDriverReactiveHealthIndicator extends AbstractReactiveHeal
 
 	@Override
 	protected Mono<Health> doHealthCheck(Health.Builder builder) {
-		return Mono.from(this.session.executeReactive(SELECT))
-				.map((row) -> builder.up().withDetail("version", row.getString(0)).build());
+		return Mono.fromSupplier(() -> {
+			Collection<Node> nodes = this.session.getMetadata().getNodes().values();
+			Optional<Node> nodeUp = nodes.stream().filter((node) -> node.getState() == NodeState.UP).findAny();
+			builder.status(nodeUp.isPresent() ? Status.UP : Status.DOWN);
+			nodeUp.map(Node::getCassandraVersion).ifPresent((version) -> builder.withDetail("version", version));
+			return builder.build();
+		});
 	}
 
 }

spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/cassandra/CassandraDriverHealthIndicatorTests.java

@@ -16,11 +16,16 @@
 
 package org.springframework.boot.actuate.cassandra;
 
+import java.util.HashMap;
+import java.util.Map;
+import java.util.UUID;
+
 import com.datastax.oss.driver.api.core.CqlSession;
 import com.datastax.oss.driver.api.core.DriverTimeoutException;
-import com.datastax.oss.driver.api.core.cql.ResultSet;
-import com.datastax.oss.driver.api.core.cql.Row;
-import com.datastax.oss.driver.api.core.cql.SimpleStatement;
+import com.datastax.oss.driver.api.core.Version;
+import com.datastax.oss.driver.api.core.metadata.Metadata;
+import com.datastax.oss.driver.api.core.metadata.Node;
+import com.datastax.oss.driver.api.core.metadata.NodeState;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.boot.actuate.health.Health;
@@ -28,14 +33,15 @@ import org.springframework.boot.actuate.health.Status;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
-import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.BDDMockito.given;
-import static org.mockito.Mockito.mock;
+import static org.mockito.BDDMockito.mock;
+import static org.mockito.BDDMockito.when;
 
 /**
  * Tests for {@link CassandraDriverHealthIndicator}.
  *
  * @author Alexandre Dutra
+ * @author Tomasz Lelek
  * @since 2.4.0
  */
 class CassandraDriverHealthIndicatorTests {
@@ -46,24 +52,137 @@ class CassandraDriverHealthIndicatorTests {
 	}
 
 	@Test
-	void healthWithCassandraUp() {
+	void oneHealthyNodeShouldReturnUp() {
+		CqlSession session = mock(CqlSession.class);
+		Metadata metadata = mock(Metadata.class);
+		Node healthyNode = mock(Node.class);
+		given(healthyNode.getState()).willReturn(NodeState.UP);
+		given(session.getMetadata()).willReturn(metadata);
+		given(metadata.getNodes()).willReturn(createNodesMap(healthyNode));
+		CassandraDriverHealthIndicator healthIndicator = new CassandraDriverHealthIndicator(session);
+		Health health = healthIndicator.health();
+		assertThat(health.getStatus()).isEqualTo(Status.UP);
+	}
+
+	@Test
+	void oneUnhealthyNodeShouldReturnDown() {
+		CqlSession session = mock(CqlSession.class);
+		Metadata metadata = mock(Metadata.class);
+		Node unhealthyNode = mock(Node.class);
+		given(unhealthyNode.getState()).willReturn(NodeState.DOWN);
+		given(session.getMetadata()).willReturn(metadata);
+		given(metadata.getNodes()).willReturn(createNodesMap(unhealthyNode));
+		CassandraDriverHealthIndicator healthIndicator = new CassandraDriverHealthIndicator(session);
+		Health health = healthIndicator.health();
+		assertThat(health.getStatus()).isEqualTo(Status.DOWN);
+	}
+
+	@Test
+	void oneUnknownNodeShouldReturnDown() {
+		CqlSession session = mock(CqlSession.class);
+		Metadata metadata = mock(Metadata.class);
+		Node unknownNode = mock(Node.class);
+		given(unknownNode.getState()).willReturn(NodeState.UNKNOWN);
+		given(session.getMetadata()).willReturn(metadata);
+		given(metadata.getNodes()).willReturn(createNodesMap(unknownNode));
+		CassandraDriverHealthIndicator healthIndicator = new CassandraDriverHealthIndicator(session);
+		Health health = healthIndicator.health();
+		assertThat(health.getStatus()).isEqualTo(Status.DOWN);
+	}
+
+	@Test
+	void oneForcedDownNodeShouldReturnDown() {
+		CqlSession session = mock(CqlSession.class);
+		Metadata metadata = mock(Metadata.class);
+		Node forcedDownNode = mock(Node.class);
+		given(forcedDownNode.getState()).willReturn(NodeState.FORCED_DOWN);
+		given(session.getMetadata()).willReturn(metadata);
+		given(metadata.getNodes()).willReturn(createNodesMap(forcedDownNode));
+		CassandraDriverHealthIndicator healthIndicator = new CassandraDriverHealthIndicator(session);
+		Health health = healthIndicator.health();
+		assertThat(health.getStatus()).isEqualTo(Status.DOWN);
+	}
+
+	@Test
+	void oneHealthyNodeAndOneUnhealthyNodeShouldReturnUp() {
+		CqlSession session = mock(CqlSession.class);
+		Metadata metadata = mock(Metadata.class);
+		Node healthyNode = mock(Node.class);
+		Node unhealthyNode = mock(Node.class);
+		given(healthyNode.getState()).willReturn(NodeState.UP);
+		given(unhealthyNode.getState()).willReturn(NodeState.DOWN);
+		given(session.getMetadata()).willReturn(metadata);
+		given(metadata.getNodes()).willReturn(createNodesMap(healthyNode, unhealthyNode));
+		CassandraDriverHealthIndicator healthIndicator = new CassandraDriverHealthIndicator(session);
+		Health health = healthIndicator.health();
+		assertThat(health.getStatus()).isEqualTo(Status.UP);
+	}
+
+	@Test
+	void oneHealthyNodeAndOneUnknownNodeShouldReturnUp() {
+		CqlSession session = mock(CqlSession.class);
+		Metadata metadata = mock(Metadata.class);
+		Node healthyNode = mock(Node.class);
+		Node unknownNode = mock(Node.class);
+		given(healthyNode.getState()).willReturn(NodeState.UP);
+		given(unknownNode.getState()).willReturn(NodeState.UNKNOWN);
+		given(session.getMetadata()).willReturn(metadata);
+		given(metadata.getNodes()).willReturn(createNodesMap(healthyNode, unknownNode));
+		CassandraDriverHealthIndicator healthIndicator = new CassandraDriverHealthIndicator(session);
+		Health health = healthIndicator.health();
+		assertThat(health.getStatus()).isEqualTo(Status.UP);
+	}
+
+	@Test
+	void oneHealthyNodeAndOneForcedDownNodeShouldReturnUp() {
 		CqlSession session = mock(CqlSession.class);
-		ResultSet resultSet = mock(ResultSet.class);
-		Row row = mock(Row.class);
-		given(session.execute(any(SimpleStatement.class))).willReturn(resultSet);
-		given(resultSet.one()).willReturn(row);
-		given(row.isNull(0)).willReturn(false);
-		given(row.getString(0)).willReturn("1.0.0");
+		Metadata metadata = mock(Metadata.class);
+		Node healthyNode = mock(Node.class);
+		Node forcedDownNode = mock(Node.class);
+		given(healthyNode.getState()).willReturn(NodeState.UP);
+		given(forcedDownNode.getState()).willReturn(NodeState.FORCED_DOWN);
+		given(session.getMetadata()).willReturn(metadata);
+		given(metadata.getNodes()).willReturn(createNodesMap(healthyNode, forcedDownNode));
 		CassandraDriverHealthIndicator healthIndicator = new CassandraDriverHealthIndicator(session);
 		Health health = healthIndicator.health();
 		assertThat(health.getStatus()).isEqualTo(Status.UP);
-		assertThat(health.getDetails().get("version")).isEqualTo("1.0.0");
+	}
+
+	@Test
+	void addVersionToDetailsIfReportedNotNull() {
+		CqlSession session = mock(CqlSession.class);
+		Metadata metadata = mock(Metadata.class);
+		when(session.getMetadata()).thenReturn(metadata);
+		Node node = mock(Node.class);
+		when(node.getState()).thenReturn(NodeState.UP);
+		when(node.getCassandraVersion()).thenReturn(Version.V4_0_0);
+		when(metadata.getNodes()).thenReturn(createNodesMap(node));
+
+		CassandraDriverHealthIndicator healthIndicator = new CassandraDriverHealthIndicator(session);
+		Health health = healthIndicator.health();
+		assertThat(health.getStatus()).isEqualTo(Status.UP);
+		assertThat(health.getDetails().get("version")).isEqualTo(Version.V4_0_0);
+	}
+
+	@Test
+	void doNotAddVersionToDetailsIfReportedNull() {
+		CqlSession session = mock(CqlSession.class);
+		Metadata metadata = mock(Metadata.class);
+		when(session.getMetadata()).thenReturn(metadata);
+		Node node = mock(Node.class);
+		when(node.getState()).thenReturn(NodeState.UP);
+		when(metadata.getNodes()).thenReturn(createNodesMap(node));
+
+		CassandraDriverHealthIndicator healthIndicator = new CassandraDriverHealthIndicator(session);
+		Health health = healthIndicator.health();
+		assertThat(health.getStatus()).isEqualTo(Status.UP);
+		assertThat(health.getDetails().get("version")).isNull();
 	}
 
 	@Test
 	void healthWithCassandraDown() {
 		CqlSession session = mock(CqlSession.class);
-		given(session.execute(any(SimpleStatement.class))).willThrow(new DriverTimeoutException("Test Exception"));
+		given(session.getMetadata()).willThrow(new DriverTimeoutException("Test Exception"));
 		CassandraDriverHealthIndicator healthIndicator = new CassandraDriverHealthIndicator(session);
 		Health health = healthIndicator.health();
 		assertThat(health.getStatus()).isEqualTo(Status.DOWN);
@@ -71,4 +190,12 @@ class CassandraDriverHealthIndicatorTests {
 				.isEqualTo(DriverTimeoutException.class.getName() + ": Test Exception");
 	}
 
+	private static Map<UUID, Node> createNodesMap(Node... nodes) {
+		Map<UUID, Node> nodesMap = new HashMap<>();
+		for (Node n : nodes) {
+			nodesMap.put(UUID.randomUUID(), n);
+		}
+		return nodesMap;
+	}
+
 }

spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/cassandra/CassandraDriverReactiveHealthIndicatorTests.java

@@ -15,15 +15,17 @@
  */
 package org.springframework.boot.actuate.cassandra;
 
-import com.datastax.dse.driver.api.core.cql.reactive.ReactiveResultSet;
-import com.datastax.dse.driver.api.core.cql.reactive.ReactiveRow;
+import java.util.HashMap;
+import java.util.Map;
+import java.util.UUID;
+
 import com.datastax.oss.driver.api.core.CqlSession;
 import com.datastax.oss.driver.api.core.DriverTimeoutException;
-import com.datastax.oss.driver.api.core.cql.SimpleStatement;
+import com.datastax.oss.driver.api.core.Version;
+import com.datastax.oss.driver.api.core.metadata.Metadata;
+import com.datastax.oss.driver.api.core.metadata.Node;
+import com.datastax.oss.driver.api.core.metadata.NodeState;
 import org.junit.jupiter.api.Test;
-import org.mockito.stubbing.Answer;
-import org.reactivestreams.Subscriber;
-import org.reactivestreams.Subscription;
 import reactor.core.publisher.Mono;
 import reactor.test.StepVerifier;
 
@@ -32,15 +34,16 @@ import org.springframework.boot.actuate.health.Status;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatIllegalArgumentException;
-import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.BDDMockito.given;
 import static org.mockito.BDDMockito.willAnswer;
 import static org.mockito.Mockito.mock;
+import static org.mockito.BDDMockito.when;
 
 /**
  * Tests for {@link CassandraDriverReactiveHealthIndicator}.
  *
  * @author Alexandre Dutra
+ * @author Tomasz Lelek
  * @since 2.4.0
  */
 class CassandraDriverReactiveHealthIndicatorTests {
@@ -51,28 +54,150 @@ class CassandraDriverReactiveHealthIndicatorTests {
 	}
 
 	@Test
-	void testCassandraIsUp() {
+	void oneHealthyNodeShouldReturnUp() {
 		CqlSession session = mock(CqlSession.class);
-		ReactiveResultSet results = mock(ReactiveResultSet.class);
-		ReactiveRow row = mock(ReactiveRow.class);
-		given(session.executeReactive(any(SimpleStatement.class))).willReturn(results);
-		willAnswer(mockReactiveResultSetBehavior(row)).given(results).subscribe(any());
-		given(row.getString(0)).willReturn("6.0.0");
-		CassandraDriverReactiveHealthIndicator cassandraReactiveHealthIndicator = new CassandraDriverReactiveHealthIndicator(
-				session);
-		Mono<Health> health = cassandraReactiveHealthIndicator.health();
+		Metadata metadata = mock(Metadata.class);
+		Node healthyNode = mock(Node.class);
+		given(healthyNode.getState()).willReturn(NodeState.UP);
+		given(session.getMetadata()).willReturn(metadata);
+		given(metadata.getNodes()).willReturn(createNodesMap(healthyNode));
+		CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session);
+		Mono<Health> health = healthIndicator.health();
+		StepVerifier.create(health).consumeNextWith((h) -> assertThat(h.getStatus()).isEqualTo(Status.UP))
+				.verifyComplete();
+	}
+
+	@Test
+	void oneUnhealthyNodeShouldReturnDown() {
+		CqlSession session = mock(CqlSession.class);
+		Metadata metadata = mock(Metadata.class);
+		Node unhealthyNode = mock(Node.class);
+		given(unhealthyNode.getState()).willReturn(NodeState.DOWN);
+		given(session.getMetadata()).willReturn(metadata);
+		given(metadata.getNodes()).willReturn(createNodesMap(unhealthyNode));
+		CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session);
+		Mono<Health> health = healthIndicator.health();
+		StepVerifier.create(health).consumeNextWith((h) -> assertThat(h.getStatus()).isEqualTo(Status.DOWN))
+				.verifyComplete();
+	}
+
+	@Test
+	void oneUnknownNodeShouldReturnDown() {
+		CqlSession session = mock(CqlSession.class);
+		Metadata metadata = mock(Metadata.class);
+		Node unknownNode = mock(Node.class);
+		given(unknownNode.getState()).willReturn(NodeState.UNKNOWN);
+		given(session.getMetadata()).willReturn(metadata);
+		given(metadata.getNodes()).willReturn(createNodesMap(unknownNode));
+		CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session);
+		Mono<Health> health = healthIndicator.health();
+		StepVerifier.create(health).consumeNextWith((h) -> assertThat(h.getStatus()).isEqualTo(Status.DOWN))
+				.verifyComplete();
+	}
+
+	@Test
+	void oneForcedDownNodeShouldReturnDown() {
+		CqlSession session = mock(CqlSession.class);
+		Metadata metadata = mock(Metadata.class);
+		Node forcedDownNode = mock(Node.class);
+		given(forcedDownNode.getState()).willReturn(NodeState.FORCED_DOWN);
+		given(session.getMetadata()).willReturn(metadata);
+		given(metadata.getNodes()).willReturn(createNodesMap(forcedDownNode));
+		CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session);
+		Mono<Health> health = healthIndicator.health();
+		StepVerifier.create(health).consumeNextWith((h) -> assertThat(h.getStatus()).isEqualTo(Status.DOWN))
+				.verifyComplete();
+	}
+
+	@Test
+	void oneHealthyNodeAndOneUnhealthyNodeShouldReturnUp() {
+		CqlSession session = mock(CqlSession.class);
+		Metadata metadata = mock(Metadata.class);
+		Node healthyNode = mock(Node.class);
+		Node unhealthyNode = mock(Node.class);
+		given(healthyNode.getState()).willReturn(NodeState.UP);
+		given(unhealthyNode.getState()).willReturn(NodeState.DOWN);
+		given(session.getMetadata()).willReturn(metadata);
+		given(metadata.getNodes()).willReturn(createNodesMap(healthyNode, unhealthyNode));
+		CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session);
+		Mono<Health> health = healthIndicator.health();
+		StepVerifier.create(health).consumeNextWith((h) -> assertThat(h.getStatus()).isEqualTo(Status.UP))
+				.verifyComplete();
+	}
+
+	@Test
+	void oneHealthyNodeAndOneUnknownNodeShouldReturnUp() {
+		CqlSession session = mock(CqlSession.class);
+		Metadata metadata = mock(Metadata.class);
+		Node healthyNode = mock(Node.class);
+		Node unknownNode = mock(Node.class);
+		given(healthyNode.getState()).willReturn(NodeState.UP);
+		given(unknownNode.getState()).willReturn(NodeState.UNKNOWN);
+		given(session.getMetadata()).willReturn(metadata);
+		given(metadata.getNodes()).willReturn(createNodesMap(healthyNode, unknownNode));
+		CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session);
+		Mono<Health> health = healthIndicator.health();
+		StepVerifier.create(health).consumeNextWith((h) -> assertThat(h.getStatus()).isEqualTo(Status.UP))
+				.verifyComplete();
+	}
+
+	@Test
+	void oneHealthyNodeAndOneForcedDownNodeShouldReturnUp() {
+		CqlSession session = mock(CqlSession.class);
+		Metadata metadata = mock(Metadata.class);
+		Node healthyNode = mock(Node.class);
+		Node forcedDownNode = mock(Node.class);
+		given(healthyNode.getState()).willReturn(NodeState.UP);
+		given(forcedDownNode.getState()).willReturn(NodeState.FORCED_DOWN);
+		given(session.getMetadata()).willReturn(metadata);
+		given(metadata.getNodes()).willReturn(createNodesMap(healthyNode, forcedDownNode));
+		CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session);
+		Mono<Health> health = healthIndicator.health();
+		StepVerifier.create(health).consumeNextWith((h) -> assertThat(h.getStatus()).isEqualTo(Status.UP))
+				.verifyComplete();
+	}
+
+	@Test
+	void addVersionToDetailsIfReportedNotNull() {
+		CqlSession session = mock(CqlSession.class);
+		Metadata metadata = mock(Metadata.class);
+		when(session.getMetadata()).thenReturn(metadata);
+		Node node = mock(Node.class);
+		when(node.getState()).thenReturn(NodeState.UP);
+		when(node.getCassandraVersion()).thenReturn(Version.V4_0_0);
+		when(metadata.getNodes()).thenReturn(createNodesMap(node));
+
+		CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session);
+		Mono<Health> health = healthIndicator.health();
 		StepVerifier.create(health).consumeNextWith((h) -> {
 			assertThat(h.getStatus()).isEqualTo(Status.UP);
 			assertThat(h.getDetails()).containsOnlyKeys("version");
-			assertThat(h.getDetails().get("version")).isEqualTo("6.0.0");
+			assertThat(h.getDetails().get("version")).isEqualTo(Version.V4_0_0);
+		}).verifyComplete();
+	}
+
+	@Test
+	void doNotAddVersionToDetailsIfReportedNull() {
+		CqlSession session = mock(CqlSession.class);
+		Metadata metadata = mock(Metadata.class);
+		when(session.getMetadata()).thenReturn(metadata);
+		Node node = mock(Node.class);
+		when(node.getState()).thenReturn(NodeState.UP);
+		when(metadata.getNodes()).thenReturn(createNodesMap(node));
+
+		CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session);
+		Mono<Health> health = healthIndicator.health();
+		StepVerifier.create(health).consumeNextWith((h) -> {
+			assertThat(h.getStatus()).isEqualTo(Status.UP);
+			assertThat(h.getDetails().get("version")).isNull();
 		}).verifyComplete();
 	}
 
 	@Test
 	void testCassandraIsDown() {
 		CqlSession session = mock(CqlSession.class);
-		given(session.executeReactive(any(SimpleStatement.class)))
-				.willThrow(new DriverTimeoutException("Test Exception"));
+		given(session.getMetadata()).willThrow(new DriverTimeoutException("Test Exception"));
+
 		CassandraDriverReactiveHealthIndicator cassandraReactiveHealthIndicator = new CassandraDriverReactiveHealthIndicator(
 				session);
 		Mono<Health> health = cassandraReactiveHealthIndicator.health();
@@ -84,23 +209,12 @@ class CassandraDriverReactiveHealthIndicatorTests {
 		}).verifyComplete();
 	}
 
-	private Answer<Void> mockReactiveResultSetBehavior(ReactiveRow row) {
-		return (invocation) -> {
-			Subscriber<ReactiveRow> subscriber = invocation.getArgument(0);
-			Subscription s = new Subscription() {
-				@Override
-				public void request(long n) {
-					subscriber.onNext(row);
-					subscriber.onComplete();
-				}
-
-				@Override
-				public void cancel() {
+	private static Map<UUID, Node> createNodesMap(Node... nodes) {
+		Map<UUID, Node> nodesMap = new HashMap<>();
+		for (Node n : nodes) {
+			nodesMap.put(UUID.randomUUID(), n);
 		}
-			};
-			subscriber.onSubscribe(s);
-			return null;
-		};
+		return nodesMap;
 	}
 
 }