content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *     http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing, software
   * distributed under the License is distributed on an "AS IS" BASIS,
   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   * See the License for the specific language governing permissions and
   * limitations under the License.
   */
  
  package org.apache.cassandra.service.reads.range;
  
  import java.util.ArrayList;
  import java.util.List;
  
  import org.apache.cassandra.schema.TableId;
  import org.junit.BeforeClass;
  import org.junit.Test;
  
  import org.apache.cassandra.SchemaLoader;
  import org.apache.cassandra.db.ConsistencyLevel;
  import org.apache.cassandra.db.Keyspace;
  import org.apache.cassandra.db.PartitionPosition;
  import org.apache.cassandra.dht.AbstractBounds;
  import org.apache.cassandra.dht.Bounds;
  import org.apache.cassandra.dht.ExcludingBounds;
  import org.apache.cassandra.dht.IncludingExcludingBounds;
  import org.apache.cassandra.dht.Range;
  import org.apache.cassandra.locator.InetAddressAndPort;
  import org.apache.cassandra.schema.KeyspaceParams;
  import org.apache.cassandra.schema.SchemaConstants;
  
  import static org.apache.cassandra.Util.rp;
  import static org.apache.cassandra.Util.token;
  import static org.junit.Assert.assertEquals;
  
  public class ReplicaPlanIteratorTest
  {
      private static final String KEYSPACE = "ReplicaPlanIteratorTest";
      private static final TableId TABLE_ID = TableId.generate();
      private static Keyspace keyspace;
  
      @BeforeClass
      public static void beforeClass() throws Throwable
      {
          SchemaLoader.prepareServer();
          SchemaLoader.createKeyspace(KEYSPACE, KeyspaceParams.simple(1));
          keyspace = Keyspace.open(KEYSPACE);
  
          new TokenUpdater().withKeys(InetAddressAndPort.getByName("127.0.0.1"), "1")
                            .withKeys(InetAddressAndPort.getByName("127.0.0.6"), "6")
                            .update();
      }
  
      @Test
      public void testRanges()
      {
          // no splits
          testRanges(range(rp("2"), rp("5")), range(rp("2"), rp("5")));
          testRanges(bounds(rp("2"), rp("5")), bounds(rp("2"), rp("5")));
          testRanges(exBounds(rp("2"), rp("5")), exBounds(rp("2"), rp("5")));
          // single split testGRR(range("2", "7"), range(rp("2"), endOf("6")), range(endOf("6"), rp("7")));
          testRanges(bounds(rp("2"), rp("7")), bounds(rp("2"), endOf("6")), range(endOf("6"), rp("7")));
          testRanges(exBounds(rp("2"), rp("7")), range(rp("2"), endOf("6")), exBounds(endOf("6"), rp("7")));
          testRanges(incExBounds(rp("2"), rp("7")), bounds(rp("2"), endOf("6")), exBounds(endOf("6"), rp("7")));
          // single split starting from min
          testRanges(range(rp(""), rp("2")), range(rp(""), endOf("1")), range(endOf("1"), rp("2")));
          testRanges(bounds(rp(""), rp("2")), bounds(rp(""), endOf("1")), range(endOf("1"), rp("2")));
          testRanges(exBounds(rp(""), rp("2")), range(rp(""), endOf("1")), exBounds(endOf("1"), rp("2")));
          testRanges(incExBounds(rp(""), rp("2")), bounds(rp(""), endOf("1")), exBounds(endOf("1"), rp("2")));
          // single split ending with max
          testRanges(range(rp("5"), rp("")), range(rp("5"), endOf("6")), range(endOf("6"), rp("")));
          testRanges(bounds(rp("5"), rp("")), bounds(rp("5"), endOf("6")), range(endOf("6"), rp("")));
          testRanges(exBounds(rp("5"), rp("")), range(rp("5"), endOf("6")), exBounds(endOf("6"), rp("")));
          testRanges(incExBounds(rp("5"), rp("")), bounds(rp("5"), endOf("6")), exBounds(endOf("6"), rp("")));
          // two splits
          testRanges(range(rp("0"), rp("7")), range(rp("0"), endOf("1")), range(endOf("1"), endOf("6")), range(endOf("6"), rp("7")));
          testRanges(bounds(rp("0"), rp("7")), bounds(rp("0"), endOf("1")), range(endOf("1"), endOf("6")), range(endOf("6"), rp("7")));
          testRanges(exBounds(rp("0"), rp("7")), range(rp("0"), endOf("1")), range(endOf("1"), endOf("6")), exBounds(endOf("6"), rp("7")));
          testRanges(incExBounds(rp("0"), rp("7")), bounds(rp("0"), endOf("1")), range(endOf("1"), endOf("6")), exBounds(endOf("6"), rp("7")));
      }
  
      @Test
      public void testExact()
      {
          // min
          testRanges(range(endOf("1"), endOf("5")), range(endOf("1"), endOf("5")));
          testRanges(range(rp("1"), endOf("5")), range(rp("1"), endOf("1")), range(endOf("1"), endOf("5")));
          testRanges(bounds(startOf("1"), endOf("5")), bounds(startOf("1"), endOf("1")), range(endOf("1"), endOf("5")));
          testRanges(exBounds(endOf("1"), rp("5")), exBounds(endOf("1"), rp("5")));
          testRanges(exBounds(rp("1"), rp("5")), range(rp("1"), endOf("1")), exBounds(endOf("1"), rp("5")));
          testRanges(exBounds(startOf("1"), endOf("5")), range(startOf("1"), endOf("1")), exBounds(endOf("1"), endOf("5")));
          testRanges(incExBounds(rp("1"), rp("5")), bounds(rp("1"), endOf("1")), exBounds(endOf("1"), rp("5")));
          // max
          testRanges(range(endOf("2"), endOf("6")), range(endOf("2"), endOf("6")));
          testRanges(bounds(startOf("2"), endOf("6")), bounds(startOf("2"), endOf("6")));
          testRanges(exBounds(rp("2"), rp("6")), exBounds(rp("2"), rp("6")));
          testRanges(incExBounds(rp("2"), rp("6")), incExBounds(rp("2"), rp("6")));
          // bothKeys
          testRanges(range(rp("1"), rp("6")), range(rp("1"), endOf("1")), range(endOf("1"), rp("6")));
          testRanges(bounds(rp("1"), rp("6")), bounds(rp("1"), endOf("1")), range(endOf("1"), rp("6")));
          testRanges(exBounds(rp("1"), rp("6")), range(rp("1"), endOf("1")), exBounds(endOf("1"), rp("6")));
          testRanges(incExBounds(rp("1"), rp("6")), bounds(rp("1"), endOf("1")), exBounds(endOf("1"), rp("6")));
      }
  
      @Test
      public void testWrapped()
      {
          // one token in wrapped range
          testRanges(range(rp("7"), rp("0")), range(rp("7"), rp("")), range(rp(""), rp("0")));
          // two tokens in wrapped range
          testRanges(range(rp("5"), rp("0")), range(rp("5"), endOf("6")), range(endOf("6"), rp("")), range(rp(""), rp("0")));
          testRanges(range(rp("7"), rp("2")), range(rp("7"), rp("")), range(rp(""), endOf("1")), range(endOf("1"), rp("2")));
          // full wraps
          testRanges(range(rp("0"), rp("0")), range(rp("0"), endOf("1")), range(endOf("1"), endOf("6")), range(endOf("6"), rp("")), range(rp(""), rp("0")));
          testRanges(range(rp(""), rp("")), range(rp(""), endOf("1")), range(endOf("1"), endOf("6")), range(endOf("6"), rp("")));
          // wrap on member tokens
          testRanges(range(rp("6"), rp("6")), range(rp("6"), endOf("6")), range(endOf("6"), rp("")), range(rp(""), endOf("1")), range(endOf("1"), rp("6")));
          testRanges(range(rp("6"), rp("1")), range(rp("6"), endOf("6")), range(endOf("6"), rp("")), range(rp(""), rp("1")));
          // end wrapped
          testRanges(range(rp("5"), rp("")), range(rp("5"), endOf("6")), range(endOf("6"), rp("")));
      }
  
      @Test
      public void testExactBounds()
      {
          // equal tokens are special cased as non-wrapping for bounds
          testRanges(bounds(rp("0"), rp("0")), bounds(rp("0"), rp("0")));
          // completely empty bounds match everything
          testRanges(bounds(rp(""), rp("")), bounds(rp(""), endOf("1")), range(endOf("1"), endOf("6")), range(endOf("6"), rp("")));
          testRanges(exBounds(rp(""), rp("")), range(rp(""), endOf("1")), range(endOf("1"), endOf("6")), exBounds(endOf("6"), rp("")));
          testRanges(incExBounds(rp(""), rp("")), bounds(rp(""), endOf("1")), range(endOf("1"), endOf("6")), exBounds(endOf("6"), rp("")));
      }
  
      @Test
      public void testLocalReplicationStrategy()
      {
          Keyspace systemKeyspace = Keyspace.open(SchemaConstants.SYSTEM_KEYSPACE_NAME);
  
          // ranges that would produce multiple splits with not-local strategy, but not with local strategy
          testRanges(systemKeyspace, range(rp("0"), rp("7")), range(rp("0"), rp("7")));
          testRanges(systemKeyspace, bounds(rp("0"), rp("7")), bounds(rp("0"), rp("7")));
          testRanges(systemKeyspace, exBounds(rp("0"), rp("7")), exBounds(rp("0"), rp("7")));
          testRanges(systemKeyspace, incExBounds(rp("0"), rp("7")), incExBounds(rp("0"), rp("7")));
  
          // wrapping ranges that should be unwrapped but not further splitted
          testRanges(systemKeyspace, range(rp("7"), rp("0")), range(rp("7"), rp("")), range(rp(""), rp("0")));
          testRanges(systemKeyspace, range(rp("7"), rp("2")), range(rp("7"), rp("")), range(rp(""), rp("2")));
      }
  
      @SafeVarargs
      private final void testRanges(AbstractBounds<PartitionPosition> queryRange, AbstractBounds<PartitionPosition>... expected)
      {
          testRanges(keyspace, queryRange, expected);
      }
  
      @SafeVarargs
      private final void testRanges(Keyspace keyspace, AbstractBounds<PartitionPosition> queryRange, AbstractBounds<PartitionPosition>... expected)
      {
          try (ReplicaPlanIterator iterator = new ReplicaPlanIterator(queryRange, null, keyspace, TABLE_ID, ConsistencyLevel.ANY))
          {
              List<AbstractBounds<PartitionPosition>> restrictedRanges = new ArrayList<>(expected.length);
              while (iterator.hasNext())
                  restrictedRanges.add(iterator.next().range());
  
              // verify range counts
              assertEquals(expected.length, restrictedRanges.size());
              assertEquals(expected.length, iterator.size());
  
              // verify the ranges
              for (int i = 0; i < expected.length; i++)
                  assertEquals("Mismatch for index " + i + ": " + restrictedRanges, expected[i], restrictedRanges.get(i));
          }
      }
  
      private static Range<PartitionPosition> range(PartitionPosition left, PartitionPosition right)
      {
          return new Range<>(left, right);
      }
  
      private static Bounds<PartitionPosition> bounds(PartitionPosition left, PartitionPosition right)
      {
          return new Bounds<>(left, right);
      }
  
      private static ExcludingBounds<PartitionPosition> exBounds(PartitionPosition left, PartitionPosition right)
      {
          return new ExcludingBounds<>(left, right);
      }
  
      private static IncludingExcludingBounds<PartitionPosition> incExBounds(PartitionPosition left, PartitionPosition right)
      {
          return new IncludingExcludingBounds<>(left, right);
      }
  
      private static PartitionPosition startOf(String key)
      {
          return token(key).minKeyBound();
      }
  
      private static PartitionPosition endOf(String key)
      {
          return token(key).maxKeyBound();
      }
  }
  