Add simple test of the .mapContinuation feature

FoundationDB · May 10, 2022 · b77a310 · b77a310
1 parent bcfed47
commit b77a310
Showing 1 changed file with 55 additions and 0 deletions.
diff --git a/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/RecordCursorTest.java b/fdb-record-layer-core/src/test/java/com/apple/foundationdb/record/RecordCursorTest.java
@@ -33,6 +33,7 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Iterators;
 import com.google.common.collect.Lists;
+import com.google.protobuf.ByteString;
 import org.apache.commons.lang3.tuple.Pair;
 import org.hamcrest.Matchers;
 import org.junit.jupiter.api.BeforeEach;
@@ -306,6 +307,60 @@ void filterTest() {
         assertEquals(Arrays.asList(2, 4, 6), newInts);
     }
 
+    @Test
+    void mapContinuationsTest() {
+        final List<Integer> ints = Arrays.asList(1, 2, 3, 4, 5, 6);
+        final ByteString prefix = ByteString.copyFromUtf8("prefix+");
+        RecordCursor<Integer> cursor = RecordCursor.fromList(ints)
+                .mapContinuation(continuation -> {
+                    if (continuation.isEnd()) {
+                        return RecordCursorEndContinuation.END;
+                    }
+
+                    return new RecordCursorContinuation() {
+                        @Nonnull
+                        @Override
+                        public ByteString toByteString() {
+                            return prefix.concat(continuation.toByteString());
+                        }
+
+                        @Nullable
+                        @Override
+                        public byte[] toBytes() {
+                            return toByteString().toByteArray();
+                        }
+
+                        @Override
+                        public boolean isEnd() {
+                            return false;
+                        }
+                    };
+                });
+
+        List<Integer> soFar = new ArrayList<>();
+        RecordCursorResult<Integer> result;
+        do {
+            result = cursor.getNext();
+            if (result.getContinuation().isEnd()) {
+                assertEquals(ByteString.EMPTY, result.getContinuation().toByteString());
+                assertNull(result.getContinuation().toBytes());
+            } else {
+                soFar.add(result.get());
+
+                // Modified continuation should begin with the prefix
+                assertTrue(result.getContinuation().toByteString().startsWith(prefix));
+
+                // Stripping away the prefix and resuming the cursor should produce the rest of the list
+                byte[] continuation = result.getContinuation().toBytes();
+                assertNotNull(continuation);
+                RecordCursor<Integer> tailCursor = RecordCursor.fromList(ints, Arrays.copyOfRange(continuation, prefix.size(), continuation.length));
+                final List<Integer> resultList = new ArrayList<>(soFar);
+                tailCursor.forEach(resultList::add).join();
+                assertEquals(ints, resultList);
+            }
+        } while (result.hasNext());
+    }
+
     @Test
     void firstTest() throws Exception {
         List<Integer> ints = Arrays.asList(1, 2, 3, 4);