Fix the tests

saikatsarkar056 · Sep 17, 2024 · c98a4d9 · c98a4d9
1 parent fba78b3
commit c98a4d9
Show file tree

Hide file tree

Showing 2 changed files with 43 additions and 2 deletions.
diff --git a/...lasticsearch/xpack/inference/external/request/ibmwatsonx/IbmWatsonxEmbeddingsRequest.java b/...lasticsearch/xpack/inference/external/request/ibmwatsonx/IbmWatsonxEmbeddingsRequest.java
@@ -50,11 +50,27 @@ public HttpRequest createHttpRequest() {
         httpPost.setEntity(byteEntity);
         httpPost.setHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaType());
 
-        IbmWatsonxRequest.decorateWithBearerToken(httpPost, model.getSecretSettings(), model.getInferenceEntityId());
+        decorateWithAuth(httpPost);
 
         return new HttpRequest(httpPost, getInferenceEntityId());
     }
 
+    public void decorateWithAuth(HttpPost httpPost) {
+        IbmWatsonxRequest.decorateWithBearerToken(httpPost, model.getSecretSettings(), model.getInferenceEntityId());
+    }
+
+    public Truncator truncator() {
+        return truncator;
+    }
+
+    public Truncator.TruncationResult truncationResult() {
+        return truncationResult;
+    }
+
+    public IbmWatsonxEmbeddingsModel model() {
+        return model;
+    }
+
     @Override
     public String getInferenceEntityId() {
         return model.getInferenceEntityId();

diff --git a/...ck/inference/external/request/ibmwatsonx/embeddings/IbmWatsonxEmbeddingsRequestTests.java b/...ck/inference/external/request/ibmwatsonx/embeddings/IbmWatsonxEmbeddingsRequestTests.java
@@ -15,7 +15,9 @@
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.inference.common.Truncator;
 import org.elasticsearch.xpack.inference.common.TruncatorTests;
+import org.elasticsearch.xpack.inference.external.request.Request;
 import org.elasticsearch.xpack.inference.external.request.ibmwatsonx.IbmWatsonxEmbeddingsRequest;
+import org.elasticsearch.xpack.inference.services.ibmwatsonx.embeddings.IbmWatsonxEmbeddingsModel;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.embeddings.IbmWatsonxEmbeddingsModelTests;
 
 import java.io.IOException;
@@ -30,6 +32,8 @@
 import static org.hamcrest.Matchers.is;
 
 public class IbmWatsonxEmbeddingsRequestTests extends ESTestCase {
+    private static final String AUTH_HEADER_VALUE = "foo";
+
     public void testCreateRequest() throws IOException {
         var model = "model";
         var projectId = "project_id";
@@ -125,10 +129,31 @@ public static IbmWatsonxEmbeddingsRequest createRequest(
     ) {
         var embeddingsModel = IbmWatsonxEmbeddingsModelTests.createModel(model, projectId, uri, apiVersion, apiKey, maxTokens, dimensions);
 
-        return new IbmWatsonxEmbeddingsRequest(
+        return new IbmWatsonxEmbeddingsWithoutAuthRequest(
             TruncatorTests.createTruncator(),
             new Truncator.TruncationResult(List.of(input), new boolean[] { false }),
             embeddingsModel
         );
     }
+
+    private static class IbmWatsonxEmbeddingsWithoutAuthRequest extends IbmWatsonxEmbeddingsRequest {
+        IbmWatsonxEmbeddingsWithoutAuthRequest(Truncator truncator, Truncator.TruncationResult input, IbmWatsonxEmbeddingsModel model) {
+            super(truncator, input, model);
+        }
+
+        @Override
+        public void decorateWithAuth(HttpPost httpPost) {
+            httpPost.setHeader(HttpHeaders.AUTHORIZATION, AUTH_HEADER_VALUE);
+        }
+
+        @Override
+        public Request truncate() {
+            IbmWatsonxEmbeddingsRequest embeddingsRequest = (IbmWatsonxEmbeddingsRequest) super.truncate();
+            return new IbmWatsonxEmbeddingsWithoutAuthRequest(
+                embeddingsRequest.truncator(),
+                embeddingsRequest.truncationResult(),
+                embeddingsRequest.model()
+            );
+        }
+    }
 }