Add unit test for using host request header

2 files changed, 21 insertions, 0 deletions

diff --git a/model-evaluation/src/test/java/ai/vespa/models/handler/HandlerTester.java b/model-evaluation/src/test/java/ai/vespa/models/handler/HandlerTester.java
index 0da7f2ed096..9d08995414b 100644
--- a/model-evaluation/src/test/java/ai/vespa/models/handler/HandlerTester.java
+++ b/model-evaluation/src/test/java/ai/vespa/models/handler/HandlerTester.java
@@ -32,9 +32,21 @@ class HandlerTester {
         assertResponse(url, Collections.emptyMap(), expectedCode, expectedResult);
     }
 
+    void assertResponse(String url, int expectedCode, String expectedResult, Map<String, String> headers) {
+        assertResponse(url, Collections.emptyMap(), expectedCode, expectedResult, headers);
+    }
+
     void assertResponse(String url, Map<String, String> properties, int expectedCode, String expectedResult) {
+        assertResponse(url, properties, expectedCode, expectedResult, Collections.emptyMap());
+    }
+
+    void assertResponse(String url, Map<String, String> properties, int expectedCode, String expectedResult, Map<String, String> headers) {
         HttpRequest getRequest = HttpRequest.createTestRequest(url, com.yahoo.jdisc.http.HttpRequest.Method.GET, null, properties);
         HttpRequest postRequest = HttpRequest.createTestRequest(url, com.yahoo.jdisc.http.HttpRequest.Method.POST, null, properties);
+        if (headers.size() > 0) {
+            headers.forEach((k,v) -> getRequest.getJDiscRequest().headers().add(k, v));
+            headers.forEach((k,v) -> postRequest.getJDiscRequest().headers().add(k, v));
+        }
         assertResponse(getRequest, expectedCode, expectedResult);
         assertResponse(postRequest, expectedCode, expectedResult);
     }
diff --git a/model-evaluation/src/test/java/ai/vespa/models/handler/ModelsEvaluationHandlerTest.java b/model-evaluation/src/test/java/ai/vespa/models/handler/ModelsEvaluationHandlerTest.java
index 63395cf4860..5b1c7a478b1 100644
--- a/model-evaluation/src/test/java/ai/vespa/models/handler/ModelsEvaluationHandlerTest.java
+++ b/model-evaluation/src/test/java/ai/vespa/models/handler/ModelsEvaluationHandlerTest.java
@@ -53,6 +53,15 @@ public class ModelsEvaluationHandlerTest {
     }
 
     @Test
+    public void testListModelsWithDifferentHost() {
+        String url = "http://localhost/model-evaluation/v1";
+        String expected =
+                "{\"mnist_softmax\":\"http://localhost:8088/model-evaluation/v1/mnist_softmax\",\"mnist_saved\":\"http://localhost:8088/model-evaluation/v1/mnist_saved\",\"mnist_softmax_saved\":\"http://localhost:8088/model-evaluation/v1/mnist_softmax_saved\",\"xgboost_2_2\":\"http://localhost:8088/model-evaluation/v1/xgboost_2_2\",\"lightgbm_regression\":\"http://localhost:8088/model-evaluation/v1/lightgbm_regression\"}";
+        handler.assertResponse(url, 200, expected, Map.of("Host", "localhost:8088"));
+    }
+
+
+    @Test
     public void testXgBoostEvaluationWithoutBindings() {
         String url = "http://localhost/model-evaluation/v1/xgboost_2_2/eval";  // only has a single function
         String expected = "{\"cells\":[{\"address\":{},\"value\":-4.376589999999999}]}";