Add sagemaker to IT

Author: prwhelan

x-pack/plugin/inference/qa/inference-service-tests/src/javaRestTest/java/org/elasticsearch/xpack/inference/InferenceGetServicesIT.java

@@ -20,7 +20,6 @@
 
 import static org.elasticsearch.xpack.inference.InferenceBaseRestTest.assertStatusOkOrCreated;
 import static org.hamcrest.Matchers.containsInAnyOrder;
-import static org.hamcrest.Matchers.equalTo;
 
 public class InferenceGetServicesIT extends BaseMockEISAuthServerTest {
 
@@ -31,13 +30,8 @@ public static void init() {
     }
 
     public void testGetServicesWithoutTaskType() throws IOException {
-        List<Object> services = getAllServices();
-        assertThat(services.size(), equalTo(24));
-
-        var providers = providers(services);
-
         assertThat(
-            providers,
+            allProviders(),
             containsInAnyOrder(
                 List.of(
                     "alibabacloud-ai-search",
@@ -69,6 +63,10 @@ public void testGetServicesWithoutTaskType() throws IOException {
         );
     }
 
+    private Iterable<String> allProviders() throws IOException {
+        return providers(getAllServices());
+    }
+
     @SuppressWarnings("unchecked")
     private Iterable<String> providers(List<Object> services) {
         return services.stream().map(service -> {
@@ -78,13 +76,8 @@ private Iterable<String> providers(List<Object> services) {
     }
 
     public void testGetServicesWithTextEmbeddingTaskType() throws IOException {
-        List<Object> services = getServices(TaskType.TEXT_EMBEDDING);
-        assertThat(services.size(), equalTo(17));
-
-        var providers = providers(services);
-
         assertThat(
-            providers,
+            providersFor(TaskType.TEXT_EMBEDDING),
             containsInAnyOrder(
                 List.of(
                     "alibabacloud-ai-search",
@@ -109,14 +102,13 @@ public void testGetServicesWithTextEmbeddingTaskType() throws IOException {
         );
     }
 
-    public void testGetServicesWithRerankTaskType() throws IOException {
-        List<Object> services = getServices(TaskType.RERANK);
-        assertThat(services.size(), equalTo(9));
-
-        var providers = providers(services);
+    private Iterable<String> providersFor(TaskType taskType) throws IOException {
+        return providers(getServices(taskType));
+    }
 
+    public void testGetServicesWithRerankTaskType() throws IOException {
         assertThat(
-            providers,
+            providersFor(TaskType.RERANK),
             containsInAnyOrder(
                 List.of(
                     "alibabacloud-ai-search",
@@ -127,20 +119,16 @@ public void testGetServicesWithRerankTaskType() throws IOException {
                     "jinaai",
                     "test_reranking_service",
                     "voyageai",
-                    "hugging_face"
+                    "hugging_face",
+                    "amazon_sagemaker"
                 ).toArray()
             )
         );
     }
 
     public void testGetServicesWithCompletionTaskType() throws IOException {
-        List<Object> services = getServices(TaskType.COMPLETION);
-        assertThat(services.size(), equalTo(14));
-
-        var providers = providers(services);
-
         assertThat(
-            providers,
+            providersFor(TaskType.COMPLETION),
             containsInAnyOrder(
                 List.of(
                     "alibabacloud-ai-search",
@@ -164,13 +152,8 @@ public void testGetServicesWithCompletionTaskType() throws IOException {
     }
 
     public void testGetServicesWithChatCompletionTaskType() throws IOException {
-        List<Object> services = getServices(TaskType.CHAT_COMPLETION);
-        assertThat(services.size(), equalTo(8));
-
-        var providers = providers(services);
-
         assertThat(
-            providers,
+            providersFor(TaskType.CHAT_COMPLETION),
             containsInAnyOrder(
                 List.of(
                     "deepseek",
@@ -187,13 +170,8 @@ public void testGetServicesWithChatCompletionTaskType() throws IOException {
     }
 
     public void testGetServicesWithSparseEmbeddingTaskType() throws IOException {
-        List<Object> services = getServices(TaskType.SPARSE_EMBEDDING);
-        assertThat(services.size(), equalTo(7));
-
-        var providers = providers(services);
-
         assertThat(
-            providers,
+            providersFor(TaskType.SPARSE_EMBEDDING),
             containsInAnyOrder(
                 List.of(
                     "alibabacloud-ai-search",
@@ -202,7 +180,8 @@ public void testGetServicesWithSparseEmbeddingTaskType() throws IOException {
                     "elasticsearch",
                     "hugging_face",
                     "streaming_completion_test_service",
-                    "test_service"
+                    "test_service",
+                    "amazon_sagemaker"
                 ).toArray()
             )
         );