From 55727779c04f3817ce1d504cc62dceccb60b11d1 Mon Sep 17 00:00:00 2001
From: Jonathan Buttner <56361221+jonathan-buttner@users.noreply.github.com>
Date: Wed, 11 Dec 2024 12:21:48 -0500
Subject: [PATCH] [ML] Fixing streaming tests locale issue (#118481)

* Fixing the string locale

* Missing a toUpper
---
 muted-tests.yml                                            | 3 ---
 .../src/main/java/org/elasticsearch/test/ESTestCase.java   | 1 +
 .../org/elasticsearch/xpack/inference/InferenceCrudIT.java | 7 +++++--
 .../mock/TestStreamingCompletionServiceExtension.java      | 5 +++--
 4 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/muted-tests.yml b/muted-tests.yml
index c0e3c217abce2..9416113770d5a 100644
--- a/muted-tests.yml
+++ b/muted-tests.yml
@@ -315,9 +315,6 @@ tests:
 - class: org.elasticsearch.xpack.test.rest.XPackRestIT
   method: test {p0=migrate/10_reindex/Test Reindex With Unsupported Mode}
   issue: https://github.com/elastic/elasticsearch/issues/118273
-- class: org.elasticsearch.xpack.inference.InferenceCrudIT
-  method: testUnifiedCompletionInference
-  issue: https://github.com/elastic/elasticsearch/issues/118405
 - class: org.elasticsearch.xpack.security.operator.OperatorPrivilegesIT
   method: testEveryActionIsEitherOperatorOnlyOrNonOperator
   issue: https://github.com/elastic/elasticsearch/issues/118220
diff --git a/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java b/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java
index 6612f0da0c43f..f678f4af22328 100644
--- a/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java
+++ b/test/framework/src/main/java/org/elasticsearch/test/ESTestCase.java
@@ -1210,6 +1210,7 @@ public static String randomAlphaOfLength(int codeUnits) {
 
     /**
      * Generate a random string containing only alphanumeric characters.
+     * <b>The locale for the string is {@link Locale#ROOT}.</b>
      * @param length the length of the string to generate
      * @return the generated string
      */
diff --git a/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceCrudIT.java b/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceCrudIT.java
index 90d4f3a8eb33b..fc593a6a8b0fa 100644
--- a/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceCrudIT.java
+++ b/x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceCrudIT.java
@@ -24,6 +24,7 @@
 import java.util.Arrays;
 import java.util.Iterator;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
@@ -471,7 +472,7 @@ public void testSupportedStream() throws Exception {
             var events = streamInferOnMockService(modelId, TaskType.COMPLETION, input);
 
             var expectedResponses = Stream.concat(
-                input.stream().map(String::toUpperCase).map(str -> "{\"completion\":[{\"delta\":\"" + str + "\"}]}"),
+                input.stream().map(s -> s.toUpperCase(Locale.ROOT)).map(str -> "{\"completion\":[{\"delta\":\"" + str + "\"}]}"),
                 Stream.of("[DONE]")
             ).iterator();
             assertThat(events.size(), equalTo((input.size() + 1) * 2));
@@ -510,7 +511,9 @@ public void testUnifiedCompletionInference() throws Exception {
     }
 
     private static Iterator<String> expectedResultsIterator(List<String> input) {
-        return Stream.concat(input.stream().map(String::toUpperCase).map(InferenceCrudIT::expectedResult), Stream.of("[DONE]")).iterator();
+        // The Locale needs to be ROOT to match what the test service is going to respond with
+        return Stream.concat(input.stream().map(s -> s.toUpperCase(Locale.ROOT)).map(InferenceCrudIT::expectedResult), Stream.of("[DONE]"))
+            .iterator();
     }
 
     private static String expectedResult(String input) {
diff --git a/x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestStreamingCompletionServiceExtension.java b/x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestStreamingCompletionServiceExtension.java
index f7a05a27354ef..80696a285fb26 100644
--- a/x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestStreamingCompletionServiceExtension.java
+++ b/x-pack/plugin/inference/qa/test-service-plugin/src/main/java/org/elasticsearch/xpack/inference/mock/TestStreamingCompletionServiceExtension.java
@@ -43,6 +43,7 @@
 import java.util.EnumSet;
 import java.util.HashMap;
 import java.util.List;
+import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.Flow;
@@ -142,7 +143,7 @@ public void unifiedCompletionInfer(
         }
 
         private StreamingChatCompletionResults makeResults(List<String> input) {
-            var responseIter = input.stream().map(String::toUpperCase).iterator();
+            var responseIter = input.stream().map(s -> s.toUpperCase(Locale.ROOT)).iterator();
             return new StreamingChatCompletionResults(subscriber -> {
                 subscriber.onSubscribe(new Flow.Subscription() {
                     @Override
@@ -173,7 +174,7 @@ private ChunkedToXContent completionChunk(String delta) {
         }
 
         private StreamingUnifiedChatCompletionResults makeUnifiedResults(UnifiedCompletionRequest request) {
-            var responseIter = request.messages().stream().map(message -> message.content().toString().toUpperCase()).iterator();
+            var responseIter = request.messages().stream().map(message -> message.content().toString().toUpperCase(Locale.ROOT)).iterator();
             return new StreamingUnifiedChatCompletionResults(subscriber -> {
                 subscriber.onSubscribe(new Flow.Subscription() {
                     @Override