[ML] Fix streaming test regex

Replace regex with string parsing. Fix elastic#114788
prwhelan · Oct 24, 2024 · f0bf4ea · f0bf4ea
1 parent 4fb7a4f
commit f0bf4ea
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 14 deletions.
diff --git a/muted-tests.yml b/muted-tests.yml
@@ -241,9 +241,6 @@ tests:
 - class: org.elasticsearch.xpack.security.operator.OperatorPrivilegesIT
   method: testEveryActionIsEitherOperatorOnlyOrNonOperator
   issue: https://github.com/elastic/elasticsearch/issues/102992
-- class: org.elasticsearch.xpack.inference.rest.ServerSentEventsRestActionListenerTests
-  method: testNoStream
-  issue: https://github.com/elastic/elasticsearch/issues/114788
 - class: org.elasticsearch.xpack.remotecluster.RemoteClusterSecurityWithApmTracingRestIT
   method: testTracingCrossCluster
   issue: https://github.com/elastic/elasticsearch/issues/112731

diff --git a/.../java/org/elasticsearch/xpack/inference/rest/ServerSentEventsRestActionListenerTests.java b/.../java/org/elasticsearch/xpack/inference/rest/ServerSentEventsRestActionListenerTests.java
@@ -64,7 +64,6 @@
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
-import java.util.regex.Pattern;
 
 import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
@@ -414,17 +413,19 @@ public void testErrorMidStream() {
         assertThat(collector.stringsVerified.getLast(), equalTo(expectedExceptionAsServerSentEvent));
     }
 
-    public void testNoStream() throws IOException {
-        var pattern = Pattern.compile("^\uFEFFevent: message\ndata: \\{\"result\":\".*\"}\n\n\uFEFFevent: message\ndata: \\[DONE]\n\n$");
+    public void testNoStream() {
+        var collector = new RandomStringCollector();
+        var expectedTestCount = randomIntBetween(2, 30);
         var request = new Request(RestRequest.Method.POST.name(), NO_STREAM_ROUTE);
-        var response = getRestClient().performRequest(request);
-        assertThat(response.getStatusLine().getStatusCode(), is(HttpStatus.SC_OK));
-        var responseString = EntityUtils.toString(response.getEntity(), StandardCharsets.UTF_8);
-
-        assertThat(
-            "Expected " + responseString + " to match pattern " + pattern.pattern(),
-            pattern.matcher(responseString).matches(),
-            is(true)
+        request.setOptions(
+            RequestOptions.DEFAULT.toBuilder()
+                .setHttpAsyncResponseConsumerFactory(() -> new AsyncResponseConsumer(collector))
+                .addParameter(REQUEST_COUNT, String.valueOf(expectedTestCount))
+                .build()
         );
+        var response = callAsync(request);
+        assertThat(response.getStatusLine().getStatusCode(), is(HttpStatus.SC_OK));
+        assertThat(collector.stringsVerified.size(), equalTo(2)); // single payload count + done byte
+        assertThat(collector.stringsVerified.peekLast(), equalTo("[DONE]"));
     }
 }