Merge pull request #2 from xinyual/addMoreUTForDelimiterChunker

add more UTs
yuye-aws · Feb 27, 2024 · e104e1b · e104e1b
2 parents 1088125 + 151162c
commit e104e1b
Showing 1 changed file with 21 additions and 14 deletions.
diff --git a/src/test/java/org/opensearch/neuralsearch/processor/chunker/DelimiterChunkerTests.java b/src/test/java/org/opensearch/neuralsearch/processor/chunker/DelimiterChunkerTests.java
@@ -40,8 +40,7 @@ public void testChunker() {
         String content = "a\nb\nc\nd";
         Map<String, Object> inputParameters = Map.of(DELIMITER_FIELD, "\n");
         List<String> chunkResult = chunker.chunk(content, inputParameters);
-        assertEquals(4, chunkResult.size());
-        assertEquals(7, cntLength(chunkResult));
+        assertEquals(List.of("a\n", "b\n", "c\n", "d"), chunkResult);
     }
 
     @Test
@@ -50,8 +49,7 @@ public void testChunkerWithDelimiterEnd() {
         String content = "a\nb\nc\nd\n";
         Map<String, Object> inputParameters = Map.of(DELIMITER_FIELD, "\n");
         List<String> chunkResult = chunker.chunk(content, inputParameters);
-        assertEquals(4, chunkResult.size());
-        assertEquals(8, cntLength(chunkResult));
+        assertEquals(List.of("a\n", "b\n", "c\n", "d\n"), chunkResult);
     }
 
     @Test
@@ -60,8 +58,7 @@ public void testChunkerWithOnlyDelimiter() {
         String content = "\n";
         Map<String, Object> inputParameters = Map.of(DELIMITER_FIELD, "\n");
         List<String> chunkResult = chunker.chunk(content, inputParameters);
-        assertEquals(1, chunkResult.size());
-        assertEquals(1, cntLength(chunkResult));
+        assertEquals(List.of("\n"), chunkResult);
     }
 
     @Test
@@ -70,15 +67,25 @@ public void testChunkerWithAllDelimiters() {
         String content = "\n\n\n";
         Map<String, Object> inputParameters = Map.of(DELIMITER_FIELD, "\n");
         List<String> chunkResult = chunker.chunk(content, inputParameters);
-        assertEquals(3, chunkResult.size());
-        assertEquals(3, cntLength(chunkResult));
+        assertEquals(List.of("\n", "\n", "\n"), chunkResult);
     }
 
-    private int cntLength(List<String> outputs) {
-        int totalLength = 0;
-        for (String output : outputs) {
-            totalLength += output.length();
-        }
-        return totalLength;
+    @Test
+    public void testChunkerWithDifferentDelimiters() {
+        DelimiterChunker chunker = new DelimiterChunker();
+        String content = "a.b.cc.d.";
+        Map<String, Object> inputParameters = Map.of(DELIMITER_FIELD, ".");
+        List<String> chunkResult = chunker.chunk(content, inputParameters);
+        assertEquals(List.of("a.", "b.", "cc.", "d."), chunkResult);
     }
+
+    @Test
+    public void testChunkerWithStringDelimter() {
+        DelimiterChunker chunker = new DelimiterChunker();
+        String content = "\n\na\n\n\n";
+        Map<String, Object> inputParameters = Map.of(DELIMITER_FIELD, "\n\n");
+        List<String> chunkResult = chunker.chunk(content, inputParameters);
+        assertEquals(List.of("\n\n", "a\n\n", "\n"), chunkResult);
+    }
+
 }