Skip to content

Commit 1f490ee

Browse files
dkropachevdkropachev
committed
RackAwarePolicy: add some unit tests for distance calculation and query
planing There are no test for it, as result policy does not work as expected. Let's have some tests for it.
1 parent ed5cd3b commit 1f490ee

File tree

1 file changed

+194
-8
lines changed

1 file changed

+194
-8
lines changed

‎driver-core/src/test/java/com/datastax/driver/core/policies/RackAwareRoundRobinPolicyTest.java‎

Lines changed: 194 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -28,27 +28,29 @@
2828
import static com.datastax.driver.core.TestUtils.nonQuietClusterCloseOptions;
2929
import static com.google.common.collect.Lists.newArrayList;
3030
import static org.mockito.Matchers.any;
31+
import static org.mockito.Mockito.mock;
3132
import static org.mockito.Mockito.spy;
33+
import static org.mockito.Mockito.when;
3234
import static org.mockito.MockitoAnnotations.initMocks;
3335

34-
import com.datastax.driver.core.Cluster;
35-
import com.datastax.driver.core.ConsistencyLevel;
36-
import com.datastax.driver.core.DataProviders;
37-
import com.datastax.driver.core.Host;
38-
import com.datastax.driver.core.MemoryAppender;
39-
import com.datastax.driver.core.QueryTracker;
40-
import com.datastax.driver.core.ScassandraCluster;
41-
import com.datastax.driver.core.Session;
36+
import com.datastax.driver.core.*;
4237
import com.datastax.driver.core.exceptions.NoHostAvailableException;
38+
import com.google.common.collect.ImmutableList;
4339
import com.google.common.collect.Lists;
40+
import java.nio.ByteBuffer;
41+
import java.util.ArrayList;
4442
import java.util.Collection;
43+
import java.util.List;
44+
import java.util.stream.Collectors;
4545
import org.apache.log4j.Level;
4646
import org.apache.log4j.Logger;
47+
import org.assertj.core.api.Assertions;
4748
import org.mockito.ArgumentCaptor;
4849
import org.mockito.Captor;
4950
import org.mockito.Mockito;
5051
import org.testng.annotations.AfterMethod;
5152
import org.testng.annotations.BeforeMethod;
53+
import org.testng.annotations.DataProvider;
5254
import org.testng.annotations.Test;
5355

5456
public class RackAwareRoundRobinPolicyTest {
@@ -60,6 +62,57 @@ public class RackAwareRoundRobinPolicyTest {
6062

6163
@Captor private ArgumentCaptor<Collection<Host>> initHostsCaptor;
6264

65+
private LoadBalancingPolicy childPolicy;
66+
private Cluster cluster;
67+
private Host host1 = mock(Host.class);
68+
private Host host2 = mock(Host.class);
69+
private Host host3 = mock(Host.class);
70+
private Host host4 = mock(Host.class);
71+
private Host host5 = mock(Host.class);
72+
private Host host6 = mock(Host.class);
73+
private ByteBuffer routingKey = ByteBuffer.wrap(new byte[] {1, 2, 3, 4});
74+
private final Statement noneLocalConsistencyStatement =
75+
new SimpleStatement("irrelevant")
76+
.setRoutingKey(routingKey)
77+
.setConsistencyLevel(ConsistencyLevel.ONE);
78+
private final Statement localConsistencyStatement =
79+
new SimpleStatement("irrelevant")
80+
.setRoutingKey(routingKey)
81+
.setConsistencyLevel(ConsistencyLevel.LOCAL_ONE);
82+
83+
@BeforeMethod(groups = "unit")
84+
public void setUpUnitTests() {
85+
CodecRegistry codecRegistry = new CodecRegistry();
86+
cluster = mock(Cluster.class);
87+
Configuration configuration = mock(Configuration.class);
88+
ProtocolOptions protocolOptions = mock(ProtocolOptions.class);
89+
Metadata metadata = mock(Metadata.class);
90+
childPolicy = mock(LoadBalancingPolicy.class);
91+
when(cluster.getConfiguration()).thenReturn(configuration);
92+
when(configuration.getCodecRegistry()).thenReturn(codecRegistry);
93+
when(configuration.getProtocolOptions()).thenReturn(protocolOptions);
94+
when(protocolOptions.getProtocolVersion()).thenReturn(ProtocolVersion.DEFAULT);
95+
when(cluster.getMetadata()).thenReturn(metadata);
96+
when(host1.isUp()).thenReturn(true);
97+
when(host1.getRack()).thenReturn("localRack");
98+
when(host1.getDatacenter()).thenReturn("localDC");
99+
when(host2.isUp()).thenReturn(true);
100+
when(host2.getRack()).thenReturn("localRack");
101+
when(host2.getDatacenter()).thenReturn("localDC");
102+
when(host3.isUp()).thenReturn(true);
103+
when(host3.getRack()).thenReturn("remoteRack");
104+
when(host3.getDatacenter()).thenReturn("localDC");
105+
when(host4.isUp()).thenReturn(true);
106+
when(host4.getRack()).thenReturn("remoteRack");
107+
when(host4.getDatacenter()).thenReturn("localDC");
108+
when(host5.isUp()).thenReturn(true);
109+
when(host5.getRack()).thenReturn("remoteRack");
110+
when(host5.getDatacenter()).thenReturn("remoteDC");
111+
when(host6.isUp()).thenReturn(true);
112+
when(host6.getRack()).thenReturn("remoteRack");
113+
when(host6.getDatacenter()).thenReturn("remoteDC");
114+
}
115+
63116
@BeforeMethod(groups = "short")
64117
public void setUp() {
65118
initMocks(this);
@@ -971,4 +1024,137 @@ public void should_use_provided_local_dc_and_warn_if_contact_points_dont_match()
9711024
sCluster.stop();
9721025
}
9731026
}
1027+
1028+
@DataProvider(name = "queryPlanTestCases")
1029+
public Object[][] queryPlanTestCases() {
1030+
return new Object[][] {
1031+
{
1032+
0,
1033+
false,
1034+
ImmutableList.of(host1, host2, host3, host4),
1035+
ImmutableList.of(host2, host1, host4, host3),
1036+
ImmutableList.of(host1, host2, host3, host4),
1037+
ImmutableList.of(host2, host1, host4, host3)
1038+
},
1039+
{
1040+
1,
1041+
false,
1042+
ImmutableList.of(host1, host2, host3, host4),
1043+
ImmutableList.of(host2, host1, host4, host3),
1044+
ImmutableList.of(host1, host2, host3, host4, host5),
1045+
ImmutableList.of(host2, host1, host4, host3, host5)
1046+
},
1047+
{
1048+
0,
1049+
true,
1050+
ImmutableList.of(host1, host2, host3, host4),
1051+
ImmutableList.of(host2, host1, host4, host3),
1052+
ImmutableList.of(host1, host2, host3, host4),
1053+
ImmutableList.of(host2, host1, host4, host3)
1054+
},
1055+
{
1056+
1,
1057+
true,
1058+
ImmutableList.of(host1, host2, host3, host4, host5),
1059+
ImmutableList.of(host2, host1, host4, host3, host5),
1060+
ImmutableList.of(host1, host2, host3, host4, host5),
1061+
ImmutableList.of(host2, host1, host4, host3, host5)
1062+
},
1063+
{
1064+
2,
1065+
true,
1066+
ImmutableList.of(host1, host2, host3, host4, host5, host6),
1067+
ImmutableList.of(host2, host1, host4, host3, host6, host5),
1068+
ImmutableList.of(host1, host2, host3, host4, host5, host6),
1069+
ImmutableList.of(host2, host1, host4, host3, host6, host5)
1070+
},
1071+
};
1072+
}
1073+
1074+
@Test(groups = "unit", dataProvider = "queryPlanTestCases")
1075+
public void should_follow_configuration_on_query_planning(
1076+
int usedHostsPerRemoteDC,
1077+
boolean allowRemoteDCsForLocalConsistencyLevel,
1078+
List<Host> queryPlanForLocalConsistencyLevel1,
1079+
List<Host> queryPlanForLocalConsistencyLevel2,
1080+
List<Host> queryPlanForNonLocalConsistencyLevel1,
1081+
List<Host> queryPlanForNonLocalConsistencyLevel2) {
1082+
RackAwareRoundRobinPolicy policy =
1083+
new RackAwareRoundRobinPolicy(
1084+
"localDC",
1085+
"localRack",
1086+
usedHostsPerRemoteDC,
1087+
allowRemoteDCsForLocalConsistencyLevel,
1088+
false,
1089+
false);
1090+
policy.init(cluster, ImmutableList.of(host1, host3, host4, host2, host5, host6));
1091+
1092+
policy.index.set(0);
1093+
ArrayList<Host> queryPlan =
1094+
Lists.newArrayList(policy.newQueryPlan("keyspace", localConsistencyStatement));
1095+
Assertions.assertThat(queryPlan)
1096+
.containsExactly(queryPlanForLocalConsistencyLevel1.toArray(new Host[0]));
1097+
queryPlan = Lists.newArrayList(policy.newQueryPlan("keyspace", localConsistencyStatement));
1098+
Assertions.assertThat(queryPlan)
1099+
.containsExactly(queryPlanForLocalConsistencyLevel2.toArray(new Host[0]));
1100+
1101+
policy.index.set(0);
1102+
queryPlan = Lists.newArrayList(policy.newQueryPlan("keyspace", noneLocalConsistencyStatement));
1103+
Assertions.assertThat(queryPlan)
1104+
.containsExactly(queryPlanForNonLocalConsistencyLevel1.toArray(new Host[0]));
1105+
queryPlan = Lists.newArrayList(policy.newQueryPlan("keyspace", noneLocalConsistencyStatement));
1106+
Assertions.assertThat(queryPlan)
1107+
.containsExactly(queryPlanForNonLocalConsistencyLevel2.toArray(new Host[0]));
1108+
}
1109+
1110+
@DataProvider(name = "distanceTestCases")
1111+
public Object[][] distanceTestCases() {
1112+
return new Object[][] {
1113+
{
1114+
0,
1115+
ImmutableList.of(
1116+
HostDistance.LOCAL,
1117+
HostDistance.LOCAL,
1118+
HostDistance.REMOTE,
1119+
HostDistance.REMOTE,
1120+
HostDistance.IGNORED,
1121+
HostDistance.IGNORED)
1122+
},
1123+
{
1124+
1,
1125+
ImmutableList.of(
1126+
HostDistance.LOCAL,
1127+
HostDistance.LOCAL,
1128+
HostDistance.REMOTE,
1129+
HostDistance.REMOTE,
1130+
HostDistance.REMOTE,
1131+
HostDistance.IGNORED)
1132+
},
1133+
{
1134+
2,
1135+
ImmutableList.of(
1136+
HostDistance.LOCAL,
1137+
HostDistance.LOCAL,
1138+
HostDistance.REMOTE,
1139+
HostDistance.REMOTE,
1140+
HostDistance.REMOTE,
1141+
HostDistance.REMOTE)
1142+
},
1143+
};
1144+
}
1145+
1146+
@Test(groups = "unit", dataProvider = "distanceTestCases")
1147+
public void should_respect_topology_on_distance(
1148+
int usedHostsPerRemoteDC, List<HostDistance> distances) {
1149+
RackAwareRoundRobinPolicy policy =
1150+
new RackAwareRoundRobinPolicy(
1151+
"localDC", "localRack", usedHostsPerRemoteDC, false, false, false);
1152+
policy.init(cluster, ImmutableList.of(host1, host3, host4, host2, host5, host6));
1153+
List<HostDistance> resultedDistances =
1154+
ImmutableList.of(host1, host2, host3, host4, host5, host6).stream()
1155+
.map(policy::distance)
1156+
.collect(Collectors.toList());
1157+
Assertions.assertThat(resultedDistances)
1158+
.containsExactly(distances.toArray(new HostDistance[0]));
1159+
}
9741160
}

0 commit comments

Comments
 (0)