From 42472856dffa6f466199a1a462b6a6bbe95dc001 Mon Sep 17 00:00:00 2001
From: Zhangmei Li <lizhangmei@baidu.com>
Date: Wed, 10 Feb 2021 20:36:02 +0800
Subject: [PATCH] add null element test case

Change-Id: I03bd304f2c153501ae285a1c28858ac73e874f3a
---
 .../hugegraph/iterator/LimitIterator.java     | 11 ++++---
 .../unit/iterator/FilterIteratorTest.java     | 29 +++++++++++++++--
 .../unit/iterator/LimitIteratorTest.java      | 31 +++++++++++++++++--
 3 files changed, 60 insertions(+), 11 deletions(-)

diff --git a/src/main/java/com/baidu/hugegraph/iterator/LimitIterator.java b/src/main/java/com/baidu/hugegraph/iterator/LimitIterator.java
index 97f02ed6..671168bb 100644
--- a/src/main/java/com/baidu/hugegraph/iterator/LimitIterator.java
+++ b/src/main/java/com/baidu/hugegraph/iterator/LimitIterator.java
@@ -41,17 +41,18 @@ protected Iterator<T> originIterator() {
     protected final boolean fetch() {
         while (this.originIterator.hasNext()) {
             T next = this.originIterator.next();
+            if (next == null) {
+                continue;
+            }
             // Do filter
             boolean reachLimit = this.filterCallback.apply(next);
             if (reachLimit) {
                 this.closeOriginIterator();
                 return false;
             }
-            if (next != null) {
-                assert this.current == none();
-                this.current = next;
-                return true;
-            }
+            assert this.current == none();
+            this.current = next;
+            return true;
         }
         return false;
     }
diff --git a/src/test/java/com/baidu/hugegraph/unit/iterator/FilterIteratorTest.java b/src/test/java/com/baidu/hugegraph/unit/iterator/FilterIteratorTest.java
index 3e62b203..0a29d017 100644
--- a/src/test/java/com/baidu/hugegraph/unit/iterator/FilterIteratorTest.java
+++ b/src/test/java/com/baidu/hugegraph/unit/iterator/FilterIteratorTest.java
@@ -41,12 +41,12 @@ public class FilterIteratorTest extends BaseUnitTest {
 
     @Test
     public void testFilter() {
-        AtomicInteger valuesCount = new AtomicInteger(0);
+        AtomicInteger callbackCount = new AtomicInteger(0);
 
         Iterator<Integer> values = DATA.iterator();
 
         Function<Integer, Boolean> filter = value -> {
-            valuesCount.incrementAndGet();
+            callbackCount.incrementAndGet();
             return (value % 2 == 0);
         };
 
@@ -57,7 +57,7 @@ public void testFilter() {
             actual.add(results.next());
         }
 
-        Assert.assertEquals(4, valuesCount.get());
+        Assert.assertEquals(4, callbackCount.get());
         Assert.assertEquals(ImmutableList.of(2, 4), actual);
     }
 
@@ -138,6 +138,29 @@ public void testNextWithMultiTimesWithoutAnyResult() {
         });
     }
 
+    @Test
+    public void testNextWithOringinIteratorReturnNullElem() {
+        List<Integer> list = new ArrayList<>();
+        list.add(1);
+        list.add(null);
+        list.add(3);
+        Iterator<Integer> vals = list.iterator();
+
+        AtomicInteger callbackCount = new AtomicInteger(0);
+
+        Iterator<Integer> results = new FilterIterator<>(vals, val -> {
+            callbackCount.incrementAndGet();
+            return true;
+        });
+
+        Assert.assertTrue(results.hasNext());
+        for (int i = 0; i < 2; i++) {
+            results.next();
+        }
+        Assert.assertFalse(results.hasNext());
+        Assert.assertEquals(2, callbackCount.get());
+    }
+
     @Test
     public void testRemove() {
         List<Integer> list = new ArrayList<>(DATA);
diff --git a/src/test/java/com/baidu/hugegraph/unit/iterator/LimitIteratorTest.java b/src/test/java/com/baidu/hugegraph/unit/iterator/LimitIteratorTest.java
index 0d61c067..8708beb9 100644
--- a/src/test/java/com/baidu/hugegraph/unit/iterator/LimitIteratorTest.java
+++ b/src/test/java/com/baidu/hugegraph/unit/iterator/LimitIteratorTest.java
@@ -41,13 +41,13 @@ public class LimitIteratorTest extends BaseUnitTest {
 
     @Test
     public void testLimit() {
-        AtomicInteger valuesCount = new AtomicInteger(0);
+        AtomicInteger callbackCount = new AtomicInteger(0);
 
         Iterator<Integer> values = DATA.iterator();
 
         int limit = 2;
         Function<Integer, Boolean> filter = value -> {
-            return valuesCount.incrementAndGet() > limit;
+            return callbackCount.incrementAndGet() > limit;
         };
 
         Iterator<Integer> results = new LimitIterator<>(values, filter);
@@ -57,7 +57,7 @@ public void testLimit() {
             actual.add(results.next());
         }
 
-        Assert.assertEquals(3, valuesCount.get());
+        Assert.assertEquals(3, callbackCount.get());
         Assert.assertEquals(ImmutableList.of(1, 2), actual);
     }
 
@@ -101,6 +101,9 @@ public void testHasNextAndNextWithMultiTimes() {
         Assert.assertThrows(NoSuchElementException.class, () -> {
             results.next();
         });
+
+        Iterator<Integer> results2 = new LimitIterator<>(vals, val -> false);
+        Assert.assertFalse(results2.hasNext());
     }
 
     @Test
@@ -138,6 +141,28 @@ public void testNextWithMultiTimesWithoutAnyResult() {
         });
     }
 
+    @Test
+    public void testNextWithOringinIteratorReturnNullElem() {
+        List<Integer> list = new ArrayList<>();
+        list.add(1);
+        list.add(null);
+        list.add(3);
+        Iterator<Integer> vals = list.iterator();
+
+        AtomicInteger callbackCount = new AtomicInteger(0);
+
+        Iterator<Integer> results = new LimitIterator<>(vals, val -> {
+            callbackCount.incrementAndGet();
+            return false;
+        });
+        Assert.assertTrue(results.hasNext());
+        for (int i = 0; i < 2; i++) {
+            results.next();
+        }
+        Assert.assertFalse(results.hasNext());
+        Assert.assertEquals(2, callbackCount.get());
+    }
+
     @Test
     public void testRemove() {
         List<Integer> list = new ArrayList<>(DATA);