Add more test cases

Signed-off-by: Andriy Redko <andriy.redko@aiven.io>
opensearch-project · Aug 28, 2024 · 5b3087a · 5b3087a
1 parent e6487d9
commit 5b3087a
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 6 deletions.
diff --git a/...rt-reactor-netty4/src/javaRestTest/java/org/opensearch/rest/ReactorNetty4StreamingIT.java b/...rt-reactor-netty4/src/javaRestTest/java/org/opensearch/rest/ReactorNetty4StreamingIT.java
@@ -85,7 +85,7 @@ public void testStreamingRequestNoBatching() throws IOException {
         assertThat(count, equalTo(5));
     }
 
-    public void testStreamingRequestOneBatch() throws IOException, InterruptedException {
+    public void testStreamingRequestOneBatchBySize() throws IOException, InterruptedException {
         final Stream<String> stream = IntStream.range(1, 6)
             .mapToObj(id -> "{ \"index\": { \"_index\": \"test-streaming\", \"_id\": \"" + id + "\" } }\n" + "{ \"name\": \"josh\" }\n");
 
@@ -126,7 +126,7 @@ public void testStreamingRequestOneBatch() throws IOException, InterruptedExcept
         assertThat(count, equalTo(5));
     }
 
-    public void testStreamingRequestManyBatches() throws IOException {
+    public void testStreamingRequestManyBatchesBySize() throws IOException {
         final Stream<String> stream = IntStream.range(1, 6)
             .mapToObj(id -> "{ \"index\": { \"_index\": \"test-streaming\", \"_id\": \"" + id + "\" } }\n" + "{ \"name\": \"josh\" }\n");
 
@@ -169,6 +169,40 @@ public void testStreamingRequestManyBatches() throws IOException {
         assertThat(count, equalTo(5));
     }
 
+    public void testStreamingRequestManyBatchesByIntervalAndSize() throws IOException {
+        final Stream<String> stream = IntStream.range(1, 6)
+            .mapToObj(id -> "{ \"index\": { \"_index\": \"test-streaming\", \"_id\": \"" + id + "\" } }\n" + "{ \"name\": \"josh\" }\n");
+
+        final Duration delay = Duration.ofSeconds(1);
+        final StreamingRequest<ByteBuffer> streamingRequest = new StreamingRequest<>(
+            "POST",
+            "/_bulk/stream",
+            Flux.fromStream(stream).delayElements(delay).map(s -> ByteBuffer.wrap(s.getBytes(StandardCharsets.UTF_8)))
+        );
+        streamingRequest.addParameter("refresh", "true");
+        streamingRequest.addParameter("batch_interval", "3s");
+        streamingRequest.addParameter("batch_size", "5");
+
+        final StreamingResponse<ByteBuffer> streamingResponse = client().streamRequest(streamingRequest);
+
+        // We don't check for a other documents here since those may appear in any of the chunks (it is very
+        // difficult to get the timing right). But at the end, the total number of the documents is being checked.
+        StepVerifier.create(Flux.from(streamingResponse.getBody()).map(b -> new String(b.array(), StandardCharsets.UTF_8)))
+            .expectNextMatches(s -> s.contains("\"result\":\"created\"") && s.contains("\"_id\":\"1\""))
+            .expectNextMatches(s -> s.contains("\"result\":\"created\"") && s.contains("\"_id\":\"5\""))
+            .expectComplete()
+            .verify();
+
+        assertThat(streamingResponse.getStatusLine().getStatusCode(), equalTo(200));
+        assertThat(streamingResponse.getWarnings(), empty());
+
+        final Request request = new Request("GET", "/test-streaming/_count");
+        final Response response = client().performRequest(request);
+        final ObjectPath objectPath = ObjectPath.createFromResponse(response);
+        final Integer count = objectPath.evaluate("count");
+        assertThat(count, equalTo(5));
+    }
+
     public void testStreamingBadRequest() throws IOException {
         final Stream<String> stream = Stream.of(
             "{ \"index\": { \"_index\": \"test-streaming\", \"_id\": \"1\" } }\n" + "{ \"name\": \"josh\"  }\n"

diff --git a/server/src/main/java/org/opensearch/rest/action/document/RestBulkStreamingAction.java b/server/src/main/java/org/opensearch/rest/action/document/RestBulkStreamingAction.java
@@ -233,11 +233,11 @@ public boolean allowsUnsafeBuffers() {
         return true;
     }
 
-    private Flux<List<HttpChunk>> createBufferedFlux(final TimeValue batch_interval, final int batch_size, StreamingRestChannel channel) {
-        if (batch_interval != null) {
-            return Flux.from(channel).bufferTimeout(batch_size, Duration.ofMillis(batch_interval.millis()));
+    private Flux<List<HttpChunk>> createBufferedFlux(final TimeValue batchInterval, final int batchSize, StreamingRestChannel channel) {
+        if (batchInterval != null) {
+            return Flux.from(channel).bufferTimeout(batchSize, Duration.ofMillis(batchInterval.millis()));
         } else {
-            return Flux.from(channel).buffer(batch_size);
+            return Flux.from(channel).buffer(batchSize);
         }
     }
 }