From 520c640ea496c6602a26520c2ec0e25f2c731981 Mon Sep 17 00:00:00 2001
From: Shane Perry <shane@lingotek.com>
Date: Mon, 4 Mar 2013 13:10:30 -0700
Subject: [PATCH] Fixed error with iterator.  Added unit test for iterator

---
 .../service/RangeSlicesIterator.java          | 15 ++-
 .../service/RangeSlicesIteratorTest.java      | 93 +++++++++++++++++++
 2 files changed, 107 insertions(+), 1 deletion(-)
 create mode 100644 core/src/test/java/me/prettyprint/cassandra/service/RangeSlicesIteratorTest.java

diff --git a/core/src/main/java/me/prettyprint/cassandra/service/RangeSlicesIterator.java b/core/src/main/java/me/prettyprint/cassandra/service/RangeSlicesIterator.java
index 41647267a..2c8012c83 100644
--- a/core/src/main/java/me/prettyprint/cassandra/service/RangeSlicesIterator.java
+++ b/core/src/main/java/me/prettyprint/cassandra/service/RangeSlicesIterator.java
@@ -34,7 +34,14 @@ public boolean hasNext() {
 			// First time through
 			iterator = Iterators.peekingIterator(query.execute().get().getList().iterator());
 		} else if (!iterator.hasNext() && rows == query.getRowCount()) {  // only need to do another query if maximum rows were retrieved
-			
+			query.setKeys(startKey, endKey);
+			iterator = Iterators.peekingIterator(query.execute().get().getList().iterator());
+			rows = 0;
+
+			if (iterator.hasNext()) {
+				// First element is startKey which was the last element on the previous query result - skip it
+				next();
+			}
 		}
 
 		while(filter != null && iterator != null && iterator.hasNext() && !filter.accept(iterator.peek())) {
@@ -62,6 +69,12 @@ public void remove() {
 		iterator.remove();
 	}
 
+	public RangeSlicesIterator<K, N, V> setFilter(SliceFilter<Row<K, N, V>> filter) {
+		this.filter = filter;
+
+		return this;
+	}
+	
 	private void refresh() {
 		query.setKeys(startKey, endKey);
 		iterator = Iterators.peekingIterator(query.execute().get().getList().iterator());
diff --git a/core/src/test/java/me/prettyprint/cassandra/service/RangeSlicesIteratorTest.java b/core/src/test/java/me/prettyprint/cassandra/service/RangeSlicesIteratorTest.java
new file mode 100644
index 000000000..9d224ceb9
--- /dev/null
+++ b/core/src/test/java/me/prettyprint/cassandra/service/RangeSlicesIteratorTest.java
@@ -0,0 +1,93 @@
+package me.prettyprint.cassandra.service;
+
+import me.prettyprint.cassandra.BaseEmbededServerSetupTest;
+import me.prettyprint.cassandra.serializers.IntegerSerializer;
+import me.prettyprint.cassandra.serializers.StringSerializer;
+import me.prettyprint.cassandra.service.template.SliceFilter;
+import me.prettyprint.hector.api.Cluster;
+import me.prettyprint.hector.api.Keyspace;
+import me.prettyprint.hector.api.beans.Row;
+import me.prettyprint.hector.api.factory.HFactory;
+import static me.prettyprint.hector.api.factory.HFactory.*;
+import me.prettyprint.hector.api.mutation.Mutator;
+import me.prettyprint.hector.api.query.RangeSlicesQuery;
+import org.junit.AfterClass;
+import static org.junit.Assert.assertEquals;
+import org.junit.Before;
+import org.junit.BeforeClass;
+import org.junit.Test;
+
+public class RangeSlicesIteratorTest extends BaseEmbededServerSetupTest {
+
+  private static final StringSerializer se = new StringSerializer();
+  private static final IntegerSerializer is = IntegerSerializer.get();
+  
+  private static final String CF = "Standard1";
+
+  private static Cluster cluster;
+  private static Keyspace keyspace;
+
+	@BeforeClass
+	public static void setupClass() {
+    cluster = getOrCreateCluster("Test Cluster", "127.0.0.1:9170");
+    keyspace = createKeyspace("Keyspace1", cluster);
+	}
+
+  @AfterClass
+  public static void teardownClass() {
+    keyspace = null;
+    cluster = null;
+  }
+
+  @Before
+  public void setupCase() {
+    // Insert 21 rows
+    Mutator<String> m = createMutator(keyspace, se);
+    for (int i = 1; i <= 21; i++) {
+      m.addInsertion("k" + i, CF, createColumn(new Integer(i), new Integer(i), is, is));
+    }
+    m.execute();
+  }
+
+  @Test
+  public void testIterator() {
+
+		RangeSlicesQuery<String, Integer, Integer> query = HFactory.createRangeSlicesQuery(keyspace, se, is, is);
+		query.setColumnFamily(CF);
+		query.setRange(null, null, false, 10);
+
+		assertKeys(query, 3, null, "k11", null);
+		assertKeys(query, 10, "k2", null, null);
+		assertKeys(query, 21, null, null, null);
+  }
+
+  @Test
+	public void testFilter() {
+		RangeSlicesQuery<String, Integer, Integer> query = HFactory.createRangeSlicesQuery(keyspace, se, is, is);
+		query.setColumnFamily(CF);
+		query.setRange(null, null, false, 10);
+
+		SliceFilter<Row<String, Integer, Integer>> filter = new SliceFilter<Row<String, Integer, Integer>>() {
+
+			@Override
+			public boolean accept(Row<String, Integer, Integer> row)
+			{
+				return Integer.parseInt(row.getKey().replaceAll("^k", "")) < 10;
+			}
+		};
+
+		assertKeys(query, 9, null, null, filter);
+	}
+	
+  private void assertKeys(RangeSlicesQuery query, int expected, String start, String end, SliceFilter<Row<String, Integer, Integer>> filter) {
+		RangeSlicesIterator<String, Integer, Integer> iterator = new RangeSlicesIterator<String, Integer, Integer>(query, start, end).setFilter(filter);
+
+    int total = 0;
+		while(iterator.hasNext()) {
+			iterator.next().getKey();
+			total ++;
+		}
+
+    assertEquals(expected, total);
+  }
+}