address review comments

trinodb · Sep 26, 2024 · 9a02d1d · 9a02d1d
1 parent 4de2870
commit 9a02d1d
Showing 1 changed file with 12 additions and 11 deletions.
diff --git a/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaMultipleBackend.java b/gateway-ha/src/test/java/io/trino/gateway/ha/TestGatewayHaMultipleBackend.java
@@ -163,7 +163,7 @@ public void testQueryDeliveryToMultipleRoutingGroups()
     public void testTrinoClusterHostCookie()
             throws Exception
     {
-        RequestBody requestBody = RequestBody.create("SELECT 1", MediaType.get("application/json; charset=utf-8"));
+        RequestBody requestBody = RequestBody.create("SELECT 1", MEDIA_TYPE);
 
         // When X-Trino-Routing-Group is set in header, query should be routed to cluster under the routing group
         Request requestWithoutCookie =
@@ -173,11 +173,12 @@ public void testTrinoClusterHostCookie()
                         .post(requestBody)
                         .addHeader("X-Trino-Routing-Group", "scheduled")
                         .build();
-        Response responseOne = httpClient.newCall(requestWithoutCookie).execute();
-        assertThat(responseOne.body().string()).contains("http://localhost:" + routerPort);
-        List<Cookie> cookies = Cookie.parseAll(responseOne.request().url(), responseOne.headers());
-        Cookie cookie = cookies.stream().filter(c -> c.name().equals("trinoClusterHost")).collect(onlyElement());
-        assertThat(cookie.value()).isEqualTo("localhost");
+        Response responseWithoutCookie = httpClient.newCall(requestWithoutCookie).execute();
+        assertThat(responseWithoutCookie.body().string()).contains("http://localhost:" + routerPort);
+        List<Cookie> cookies = Cookie.parseAll(responseWithoutCookie.request().url(), responseWithoutCookie.headers());
+        Cookie clusterHostCookie = cookies.stream().filter(c -> c.name().equals("trinoClusterHost")).collect(onlyElement());
+        assertThat(clusterHostCookie.value()).isEqualTo("localhost");
+
         // test with sending the request which includes trinoClusterHost in the cookie
         // when X-Trino-Routing-Group is set in header, query should be routed to cluster under the routing group
         Request requestWithCookie =
@@ -188,11 +189,11 @@ public void testTrinoClusterHostCookie()
                         .addHeader("X-Trino-Routing-Group", "scheduled")
                         .addHeader("Cookie", "trinoClientHost=foo.example.com")
                         .build();
-        Response responseTwo = httpClient.newCall(requestWithCookie).execute();
-        assertThat(responseTwo.body().string()).contains("http://localhost:" + routerPort);
-        cookies = Cookie.parseAll(responseTwo.request().url(), responseTwo.headers());
-        cookie = cookies.stream().filter(c -> c.name().equals("trinoClusterHost")).collect(onlyElement());
-        assertThat(cookie.value()).isEqualTo("localhost");
+        Response responseWithCookie = httpClient.newCall(requestWithCookie).execute();
+        assertThat(responseWithCookie.body().string()).contains("http://localhost:" + routerPort);
+        List<Cookie> overridenCookies = Cookie.parseAll(responseWithCookie.request().url(), responseWithCookie.headers());
+        Cookie overridenClusterHostCookie = overridenCookies.stream().filter(c -> c.name().equals("trinoClusterHost")).collect(onlyElement());
+        assertThat(overridenClusterHostCookie.value()).isEqualTo("localhost");
     }
 
     @Test