feat: improved tests

epam · Jan 13, 2025 · caf8ed0 · caf8ed0
1 parent f6972e7
commit caf8ed0
Show file tree

Hide file tree

Showing 2 changed files with 39 additions and 3 deletions.
diff --git a/server/src/main/java/com/epam/aidial/core/server/controller/DeploymentFeatureController.java b/server/src/main/java/com/epam/aidial/core/server/controller/DeploymentFeatureController.java
@@ -76,8 +76,13 @@ private void handleRequestBody(String endpoint, boolean requireEndpoint, Buffer
     }
 
     private void handleError(Throwable error) {
+<<<<<<< Updated upstream
         log.info("Error occurred while processing request", error);
         respond(HttpStatus.INTERNAL_SERVER_ERROR, error.getMessage());
+=======
+        log.error("Error occurred while processing request", error);
+        context.respond(HttpStatus.INTERNAL_SERVER_ERROR, error.getMessage());
+>>>>>>> Stashed changes
     }
 
     @SneakyThrows

diff --git a/server/src/test/java/com/epam/aidial/core/server/FeaturesApiTest.java b/server/src/test/java/com/epam/aidial/core/server/FeaturesApiTest.java
@@ -2,12 +2,37 @@
 
 import io.vertx.core.http.HttpMethod;
 import lombok.SneakyThrows;
+import okhttp3.Headers;
+
 import org.junit.jupiter.api.Test;
 
 import java.net.URI;
+import java.util.List;
+import java.util.Map;
+import java.util.stream.StreamSupport;
+import java.util.stream.Stream;
 
 public class FeaturesApiTest extends ResourceBaseTest {
 
+    private static String[] convertHeadersToFlatArray(Headers headers) {
+        return StreamSupport.stream(headers.spliterator(), false)
+                .flatMap(header -> Stream.of(header.getFirst(), header.getSecond()))
+                .toArray(String[]::new);
+    }
+
+    private static Headers filterHeaders(Headers headers, Headers mask) {
+        Headers.Builder filteredHeaders = new Headers.Builder();
+        for (Map.Entry<String, List<String>> entry : headers.toMultimap().entrySet()) {
+            String key = entry.getKey();
+            if (mask.names().contains(key.toLowerCase())) {
+                for (String value : entry.getValue()) {
+                    filteredHeaders.add(key, value);
+                }
+            }
+        }
+        return filteredHeaders.build();
+    }
+
     @Test
     void testRateEndpointModel() {
         String inboundPath = "/v1/chat-gpt-35-turbo/rate";
@@ -64,12 +89,18 @@ void testUpstreamEndpoint(String inboundPath, String upstream) {
 
     @SneakyThrows
     void testUpstreamEndpoint(String inboundPath, String upstream, HttpMethod method) {
+        Headers requestExtraHeaders = new Headers.Builder().add("foo", "bar").build();
+        String[] requestExtraHeadersArray = convertHeadersToFlatArray(requestExtraHeaders);
+
         URI uri = URI.create(upstream);
         try (TestWebServer server = new TestWebServer(uri.getPort())) {
-            server.map(method, uri.getPath(),  request -> TestWebServer.createResponse(200, "PONG", "foo", "bar"));
+            server.map(method, uri.getPath(), request -> {
+                Headers responseHeaders = filterHeaders(request.getHeaders(), requestExtraHeaders);
+                return TestWebServer.createResponse(200, "PONG", convertHeadersToFlatArray(responseHeaders));
+            });
 
-            Response response = send(method, inboundPath, null, "", "foo", "bar");
-            verify(response, 200, "PONG", "foo", "bar");
+            Response response = send(method, inboundPath, null, "PING", requestExtraHeadersArray);
+            verify(response, 200, "PONG", requestExtraHeadersArray);
         }
     }
 }