From 8e5a041bb48b18ec4f7518a76090d80bb0a07029 Mon Sep 17 00:00:00 2001 From: tomekl007 Date: Fri, 21 Aug 2020 11:16:17 +0200 Subject: [PATCH 1/2] Improve Cassandra health indicator with more robust mechanism See gh-23041 --- .../CassandraDriverHealthIndicator.java | 22 +-- ...assandraDriverReactiveHealthIndicator.java | 21 +- .../CassandraDriverHealthIndicatorTests.java | 155 +++++++++++++-- ...draDriverReactiveHealthIndicatorTests.java | 186 ++++++++++++++---- 4 files changed, 316 insertions(+), 68 deletions(-) diff --git a/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/cassandra/CassandraDriverHealthIndicator.java b/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/cassandra/CassandraDriverHealthIndicator.java index bbb1468c0e..ed42a0bbda 100644 --- a/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/cassandra/CassandraDriverHealthIndicator.java +++ b/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 nodes = this.session.getMetadata().getNodes().values(); + Optional 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)); } } diff --git a/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/cassandra/CassandraDriverReactiveHealthIndicator.java b/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/cassandra/CassandraDriverReactiveHealthIndicator.java index 40483ca04e..893faf5855 100644 --- a/spring-boot-project/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/cassandra/CassandraDriverReactiveHealthIndicator.java +++ b/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 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 nodes = this.session.getMetadata().getNodes().values(); + Optional 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(); + }); } } diff --git a/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/cassandra/CassandraDriverHealthIndicatorTests.java b/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/cassandra/CassandraDriverHealthIndicatorTests.java index 875edd64c3..c7fa024758 100644 --- a/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/cassandra/CassandraDriverHealthIndicatorTests.java +++ b/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); - 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); + 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); - assertThat(health.getDetails().get("version")).isEqualTo("1.0.0"); + } + + @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); + 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); + } + + @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 createNodesMap(Node... nodes) { + Map nodesMap = new HashMap<>(); + for (Node n : nodes) { + nodesMap.put(UUID.randomUUID(), n); + } + return nodesMap; + } + } diff --git a/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/cassandra/CassandraDriverReactiveHealthIndicatorTests.java b/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/cassandra/CassandraDriverReactiveHealthIndicatorTests.java index 50cee9f37d..9d1066bda8 100644 --- a/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/cassandra/CassandraDriverReactiveHealthIndicatorTests.java +++ b/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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = 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 = cassandraReactiveHealthIndicator.health(); @@ -84,23 +209,12 @@ class CassandraDriverReactiveHealthIndicatorTests { }).verifyComplete(); } - private Answer mockReactiveResultSetBehavior(ReactiveRow row) { - return (invocation) -> { - Subscriber subscriber = invocation.getArgument(0); - Subscription s = new Subscription() { - @Override - public void request(long n) { - subscriber.onNext(row); - subscriber.onComplete(); - } - - @Override - public void cancel() { - } - }; - subscriber.onSubscribe(s); - return null; - }; + private static Map createNodesMap(Node... nodes) { + Map nodesMap = new HashMap<>(); + for (Node n : nodes) { + nodesMap.put(UUID.randomUUID(), n); + } + return nodesMap; } } From 6f08e9709ebc9d9f21927f036155f6ebf8846ecf Mon Sep 17 00:00:00 2001 From: Stephane Nicoll Date: Fri, 2 Oct 2020 15:47:02 +0200 Subject: [PATCH 2/2] Polish "Improve Cassandra health indicator with more robust mechanism" See gh-23041 --- .../CassandraDriverHealthIndicatorTests.java | 123 +++++++----------- ...draDriverReactiveHealthIndicatorTests.java | 123 ++++++------------ 2 files changed, 87 insertions(+), 159 deletions(-) diff --git a/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/cassandra/CassandraDriverHealthIndicatorTests.java b/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/cassandra/CassandraDriverHealthIndicatorTests.java index c7fa024758..f9d8827c8b 100644 --- a/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/cassandra/CassandraDriverHealthIndicatorTests.java +++ b/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/cassandra/CassandraDriverHealthIndicatorTests.java @@ -16,7 +16,10 @@ package org.springframework.boot.actuate.cassandra; +import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.UUID; @@ -34,15 +37,13 @@ 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.BDDMockito.given; -import static org.mockito.BDDMockito.mock; -import static org.mockito.BDDMockito.when; +import static org.mockito.Mockito.mock; /** * Tests for {@link CassandraDriverHealthIndicator}. * * @author Alexandre Dutra - * @author Tomasz Lelek - * @since 2.4.0 + * @author Stephane Nicoll */ class CassandraDriverHealthIndicatorTests { @@ -52,112 +53,70 @@ class CassandraDriverHealthIndicatorTests { } @Test - 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)); + void healthWithOneHealthyNodeShouldReturnUp() { + CqlSession session = mockCqlSessionWithNodeState(NodeState.UP); 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)); + void healthWithOneUnhealthyNodeShouldReturnDown() { + CqlSession session = mockCqlSessionWithNodeState(NodeState.DOWN); 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)); + void healthWithOneUnknownNodeShouldReturnDown() { + CqlSession session = mockCqlSessionWithNodeState(NodeState.UNKNOWN); 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)); + void healthWithOneForcedDownNodeShouldReturnDown() { + CqlSession session = mockCqlSessionWithNodeState(NodeState.FORCED_DOWN); 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)); + void healthWithOneHealthyNodeAndOneUnhealthyNodeShouldReturnUp() { + CqlSession session = mockCqlSessionWithNodeState(NodeState.UP, NodeState.DOWN); 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)); + void healthWithOneHealthyNodeAndOneUnknownNodeShouldReturnUp() { + CqlSession session = mockCqlSessionWithNodeState(NodeState.UP, NodeState.UNKNOWN); CassandraDriverHealthIndicator healthIndicator = new CassandraDriverHealthIndicator(session); Health health = healthIndicator.health(); assertThat(health.getStatus()).isEqualTo(Status.UP); } @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)); + void healthWithOneHealthyNodeAndOneForcedDownNodeShouldReturnUp() { + CqlSession session = mockCqlSessionWithNodeState(NodeState.UP, NodeState.FORCED_DOWN); CassandraDriverHealthIndicator healthIndicator = new CassandraDriverHealthIndicator(session); Health health = healthIndicator.health(); assertThat(health.getStatus()).isEqualTo(Status.UP); } @Test - void addVersionToDetailsIfReportedNotNull() { + void healthWithNodeVersionShouldAddVersionDetail() { CqlSession session = mock(CqlSession.class); Metadata metadata = mock(Metadata.class); - when(session.getMetadata()).thenReturn(metadata); + given(session.getMetadata()).willReturn(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)); - + given(node.getState()).willReturn(NodeState.UP); + given(node.getCassandraVersion()).willReturn(Version.V4_0_0); + given(metadata.getNodes()).willReturn(createNodesWithRandomUUID(Collections.singletonList(node))); CassandraDriverHealthIndicator healthIndicator = new CassandraDriverHealthIndicator(session); Health health = healthIndicator.health(); assertThat(health.getStatus()).isEqualTo(Status.UP); @@ -165,14 +124,8 @@ class CassandraDriverHealthIndicatorTests { } @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)); - + void healthWithoutNodeVersionShouldNotAddVersionDetail() { + CqlSession session = mockCqlSessionWithNodeState(NodeState.UP); CassandraDriverHealthIndicator healthIndicator = new CassandraDriverHealthIndicator(session); Health health = healthIndicator.health(); assertThat(health.getStatus()).isEqualTo(Status.UP); @@ -180,7 +133,7 @@ class CassandraDriverHealthIndicatorTests { } @Test - void healthWithCassandraDown() { + void healthWithcassandraDownShouldReturnDown() { CqlSession session = mock(CqlSession.class); given(session.getMetadata()).willThrow(new DriverTimeoutException("Test Exception")); CassandraDriverHealthIndicator healthIndicator = new CassandraDriverHealthIndicator(session); @@ -190,12 +143,24 @@ class CassandraDriverHealthIndicatorTests { .isEqualTo(DriverTimeoutException.class.getName() + ": Test Exception"); } - private static Map createNodesMap(Node... nodes) { - Map nodesMap = new HashMap<>(); - for (Node n : nodes) { - nodesMap.put(UUID.randomUUID(), n); + private CqlSession mockCqlSessionWithNodeState(NodeState... nodeStates) { + CqlSession session = mock(CqlSession.class); + Metadata metadata = mock(Metadata.class); + List nodes = new ArrayList<>(); + for (NodeState nodeState : nodeStates) { + Node node = mock(Node.class); + given(node.getState()).willReturn(nodeState); + nodes.add(node); } - return nodesMap; + given(session.getMetadata()).willReturn(metadata); + given(metadata.getNodes()).willReturn(createNodesWithRandomUUID(nodes)); + return session; + } + + private Map createNodesWithRandomUUID(List nodes) { + Map indexedNodes = new HashMap<>(); + nodes.forEach((node) -> indexedNodes.put(UUID.randomUUID(), node)); + return indexedNodes; } } diff --git a/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/cassandra/CassandraDriverReactiveHealthIndicatorTests.java b/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/cassandra/CassandraDriverReactiveHealthIndicatorTests.java index 9d1066bda8..c35deec14c 100644 --- a/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/cassandra/CassandraDriverReactiveHealthIndicatorTests.java +++ b/spring-boot-project/spring-boot-actuator/src/test/java/org/springframework/boot/actuate/cassandra/CassandraDriverReactiveHealthIndicatorTests.java @@ -15,7 +15,10 @@ */ package org.springframework.boot.actuate.cassandra; +import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; +import java.util.List; import java.util.Map; import java.util.UUID; @@ -35,16 +38,13 @@ 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.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 + * @author Stephane Nicoll */ class CassandraDriverReactiveHealthIndicatorTests { @@ -54,13 +54,8 @@ class CassandraDriverReactiveHealthIndicatorTests { } @Test - 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)); + void healthWithOneHealthyNodeShouldReturnUp() { + CqlSession session = mockCqlSessionWithNodeState(NodeState.UP); CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session); Mono health = healthIndicator.health(); StepVerifier.create(health).consumeNextWith((h) -> assertThat(h.getStatus()).isEqualTo(Status.UP)) @@ -68,13 +63,8 @@ class CassandraDriverReactiveHealthIndicatorTests { } @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)); + void healthWithOneUnhealthyNodeShouldReturnDown() { + CqlSession session = mockCqlSessionWithNodeState(NodeState.DOWN); CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session); Mono health = healthIndicator.health(); StepVerifier.create(health).consumeNextWith((h) -> assertThat(h.getStatus()).isEqualTo(Status.DOWN)) @@ -82,13 +72,8 @@ class CassandraDriverReactiveHealthIndicatorTests { } @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)); + void healthWithOneUnknownNodeShouldReturnDown() { + CqlSession session = mockCqlSessionWithNodeState(NodeState.UNKNOWN); CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session); Mono health = healthIndicator.health(); StepVerifier.create(health).consumeNextWith((h) -> assertThat(h.getStatus()).isEqualTo(Status.DOWN)) @@ -96,13 +81,8 @@ class CassandraDriverReactiveHealthIndicatorTests { } @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)); + void healthWithOneForcedDownNodeShouldReturnDown() { + CqlSession session = mockCqlSessionWithNodeState(NodeState.FORCED_DOWN); CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session); Mono health = healthIndicator.health(); StepVerifier.create(health).consumeNextWith((h) -> assertThat(h.getStatus()).isEqualTo(Status.DOWN)) @@ -110,15 +90,8 @@ class CassandraDriverReactiveHealthIndicatorTests { } @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)); + void healthWithOneHealthyNodeAndOneUnhealthyNodeShouldReturnUp() { + CqlSession session = mockCqlSessionWithNodeState(NodeState.UP, NodeState.DOWN); CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session); Mono health = healthIndicator.health(); StepVerifier.create(health).consumeNextWith((h) -> assertThat(h.getStatus()).isEqualTo(Status.UP)) @@ -126,15 +99,8 @@ class CassandraDriverReactiveHealthIndicatorTests { } @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)); + void healthWithOneHealthyNodeAndOneUnknownNodeShouldReturnUp() { + CqlSession session = mockCqlSessionWithNodeState(NodeState.UP, NodeState.UNKNOWN); CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session); Mono health = healthIndicator.health(); StepVerifier.create(health).consumeNextWith((h) -> assertThat(h.getStatus()).isEqualTo(Status.UP)) @@ -142,15 +108,8 @@ class CassandraDriverReactiveHealthIndicatorTests { } @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)); + void healthWithOneHealthyNodeAndOneForcedDownNodeShouldReturnUp() { + CqlSession session = mockCqlSessionWithNodeState(NodeState.UP, NodeState.FORCED_DOWN); CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session); Mono health = healthIndicator.health(); StepVerifier.create(health).consumeNextWith((h) -> assertThat(h.getStatus()).isEqualTo(Status.UP)) @@ -158,15 +117,14 @@ class CassandraDriverReactiveHealthIndicatorTests { } @Test - void addVersionToDetailsIfReportedNotNull() { + void healthWithNodeVersionShouldAddVersionDetail() { CqlSession session = mock(CqlSession.class); Metadata metadata = mock(Metadata.class); - when(session.getMetadata()).thenReturn(metadata); + given(session.getMetadata()).willReturn(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)); - + given(node.getState()).willReturn(NodeState.UP); + given(node.getCassandraVersion()).willReturn(Version.V4_0_0); + given(metadata.getNodes()).willReturn(createNodesWithRandomUUID(Collections.singletonList(node))); CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session); Mono health = healthIndicator.health(); StepVerifier.create(health).consumeNextWith((h) -> { @@ -177,14 +135,8 @@ class CassandraDriverReactiveHealthIndicatorTests { } @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)); - + void healthWithoutNodeVersionShouldNotAddVersionDetail() { + CqlSession session = mockCqlSessionWithNodeState(NodeState.UP); CassandraDriverReactiveHealthIndicator healthIndicator = new CassandraDriverReactiveHealthIndicator(session); Mono health = healthIndicator.health(); StepVerifier.create(health).consumeNextWith((h) -> { @@ -194,10 +146,9 @@ class CassandraDriverReactiveHealthIndicatorTests { } @Test - void testCassandraIsDown() { + void healthWithCassandraDownShouldReturnDown() { CqlSession session = mock(CqlSession.class); given(session.getMetadata()).willThrow(new DriverTimeoutException("Test Exception")); - CassandraDriverReactiveHealthIndicator cassandraReactiveHealthIndicator = new CassandraDriverReactiveHealthIndicator( session); Mono health = cassandraReactiveHealthIndicator.health(); @@ -209,12 +160,24 @@ class CassandraDriverReactiveHealthIndicatorTests { }).verifyComplete(); } - private static Map createNodesMap(Node... nodes) { - Map nodesMap = new HashMap<>(); - for (Node n : nodes) { - nodesMap.put(UUID.randomUUID(), n); + private CqlSession mockCqlSessionWithNodeState(NodeState... nodeStates) { + CqlSession session = mock(CqlSession.class); + Metadata metadata = mock(Metadata.class); + List nodes = new ArrayList<>(); + for (NodeState nodeState : nodeStates) { + Node node = mock(Node.class); + given(node.getState()).willReturn(nodeState); + nodes.add(node); } - return nodesMap; + given(session.getMetadata()).willReturn(metadata); + given(metadata.getNodes()).willReturn(createNodesWithRandomUUID(nodes)); + return session; + } + + private Map createNodesWithRandomUUID(List nodes) { + Map indexedNodes = new HashMap<>(); + nodes.forEach((node) -> indexedNodes.put(UUID.randomUUID(), node)); + return indexedNodes; } }