io.stargate.it.cql.TokenAwareTest Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of testing Show documentation
There is a newer version: 2.1.0-BETA-19
package io.stargate.it.cql;

import static org.assertj.core.api.Assertions.assertThat;

import com.datastax.oss.driver.api.core.CqlIdentifier;
import com.datastax.oss.driver.api.core.CqlSession;
import com.datastax.oss.driver.api.core.ProtocolVersion;
import com.datastax.oss.driver.api.core.metadata.Node;
import com.datastax.oss.driver.api.core.metadata.TokenMap;
import com.datastax.oss.driver.api.core.type.codec.TypeCodecs;
import io.stargate.it.BaseIntegrationTest;
import io.stargate.it.TestOrder;
import io.stargate.it.driver.CqlSessionExtension;
import io.stargate.it.driver.TestKeyspace;
import io.stargate.it.driver.WithProtocolVersion;
import io.stargate.it.storage.StargateEnvironmentInfo;
import io.stargate.it.storage.StargateSpec;
import java.util.Collection;
import java.util.HashSet;
import java.util.Set;
import java.util.UUID;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Order;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.ExtendWith;

@StargateSpec(nodes = 3)
@ExtendWith(CqlSessionExtension.class)
@Order(TestOrder.LAST)
public abstract class TokenAwareTest extends BaseIntegrationTest {
  @Test
  @DisplayName("Should use all Stargate addresses when using token-aware load balancing")
  public void testTokenMapDistribution(
      CqlSession session, StargateEnvironmentInfo stargate, @TestKeyspace CqlIdentifier keyspace) {
    assertThat(session.getMetadata().getTokenMap()).isPresent();
    final TokenMap tokenMap = session.getMetadata().getTokenMap().get();

    Collection expectedReplicasTried = session.getMetadata().getNodes().values();
    assertThat(stargate.nodes()).hasSize(expectedReplicasTried.size());

    Set allReplicasTried = new HashSet<>();
    for (int i = 0; i < 10 * expectedReplicasTried.size(); ++i) {
      Set replicas =
          tokenMap.getReplicas(
              keyspace, TypeCodecs.UUID.encode(UUID.randomUUID(), ProtocolVersion.DEFAULT));
      assertThat(replicas).isNotEmpty();
      Node replica = replicas.iterator().next();
      allReplicasTried.add(replica);
    }

    assertThat(allReplicasTried).containsExactlyInAnyOrderElementsOf(expectedReplicasTried);
  }

  @WithProtocolVersion("V4")
  public static class WithV4ProtocolVersionTest extends TokenAwareTest {}

  @WithProtocolVersion("V5")
  public static class WithV5ProtocolVersionTest extends TokenAwareTest {}
}