[Inference API] Improve completion response entity tests (#108512)

elastic · May 10, 2024 · 5a824c5 · 5a824c5
1 parent d6cb12e
commit 5a824c5
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 9 deletions.
diff --git a/...ack/inference/external/response/azureopenai/AzureOpenAiCompletionResponseEntityTests.java b/...ack/inference/external/response/azureopenai/AzureOpenAiCompletionResponseEntityTests.java
@@ -17,7 +17,6 @@
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 
-import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
 import static org.mockito.Mockito.mock;
 
@@ -50,7 +49,7 @@ public void testFromResponse_CreatesResultsForASingleItem() throws IOException {
                          "index": 0,
                          "logprobs": null,
                          "message": {
-                             "content": "response",
+                             "content": "result",
                              "role": "assistant"
                          }
                      }
@@ -92,10 +91,8 @@ public void testFromResponse_CreatesResultsForASingleItem() throws IOException {
             new HttpResult(mock(HttpResponse.class), responseJson.getBytes(StandardCharsets.UTF_8))
         );
 
-        assertThat(chatCompletionResults.getResults().size(), equalTo(1));
-
-        ChatCompletionResults.Result result = chatCompletionResults.getResults().get(0);
-        assertThat(result.asMap().get(result.getResultsField()), is("response"));
+        assertThat(chatCompletionResults.getResults().size(), is(1));
+        assertThat(chatCompletionResults.getResults().get(0).content(), is("result"));
     }
 
     public void testFromResponse_FailsWhenChoicesFieldIsNotPresent() {

diff --git a/...rch/xpack/inference/external/response/openai/OpenAiChatCompletionResponseEntityTests.java b/...rch/xpack/inference/external/response/openai/OpenAiChatCompletionResponseEntityTests.java
@@ -17,7 +17,6 @@
 import java.io.IOException;
 import java.nio.charset.StandardCharsets;
 
-import static org.hamcrest.Matchers.equalTo;
 import static org.hamcrest.Matchers.is;
 import static org.mockito.Mockito.mock;
 
@@ -35,7 +34,7 @@ public void testFromResponse_CreatesResultsForASingleItem() throws IOException {
                   "index": 0,
                   "message": {
                     "role": "assistant",
-                    "content": "some content"
+                    "content": "result"
                   },
                   "logprobs": null,
                   "finish_reason": "stop"
@@ -55,7 +54,8 @@ public void testFromResponse_CreatesResultsForASingleItem() throws IOException {
             new HttpResult(mock(HttpResponse.class), responseJson.getBytes(StandardCharsets.UTF_8))
         );
 
-        assertThat(chatCompletionResults.getResults().size(), equalTo(1));
+        assertThat(chatCompletionResults.getResults().size(), is(1));
+        assertThat(chatCompletionResults.getResults().get(0).content(), is("result"));
     }
 
     public void testFromResponse_FailsWhenChoicesFieldIsNotPresent() {