Skip to content

Commit

Permalink
Replace MockWebServer with dockerized Trino
Browse files Browse the repository at this point in the history
  • Loading branch information
@ebyhr
ebyhr committed Jan 9, 2024
1 parent 7e30bbe commit f366966
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 117 deletions.
6 changes: 0 additions & 6 deletions gateway-ha/pom.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -195,12 +195,6 @@
<scope>test</scope>
</dependency>

<dependency>
<groupId>com.squareup.okhttp</groupId>
<artifactId>mockwebserver</artifactId>
<scope>test</scope>
</dependency>

<dependency>
<groupId>org.assertj</groupId>
<artifactId>assertj-core</artifactId>
Expand Down
141 changes: 37 additions & 104 deletions gateway-ha/src/test/java/io/trino/gateway/proxyserver/TestProxyServer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,126 +13,59 @@
*/
package io.trino.gateway.proxyserver;

import com.squareup.okhttp.mockwebserver.MockResponse;
import com.squareup.okhttp.mockwebserver.MockWebServer;
import com.squareup.okhttp.mockwebserver.RecordedRequest;
import org.apache.http.HttpResponse;
import org.apache.http.client.methods.HttpGet;
import org.apache.http.client.methods.HttpUriRequest;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClientBuilder;
import org.apache.http.util.EntityUtils;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.TestInstance;
import org.testcontainers.containers.TrinoContainer;

import java.io.IOException;
import java.util.Random;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.Statement;

import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.api.TestInstance.Lifecycle.PER_CLASS;

@TestInstance(PER_CLASS)
public class TestProxyServer
{
private static int serverPort;
private static MockWebServer backend;
private static ProxyServer proxyServer;
private TrinoContainer trino;
private ProxyServer proxy;
private int proxyPort;

@Test
public void testProxyServer()
throws IOException
@BeforeEach
void setUp()
{
String mockResponseText = "Test1234";
setProxyServer(mockResponseText);
try {
proxyServer.start();
CloseableHttpClient httpclient = HttpClientBuilder.create().build();
HttpUriRequest httpUriRequest = new HttpGet("http://localhost:" + serverPort);
HttpResponse response = httpclient.execute(httpUriRequest);
assertEquals(mockResponseText, EntityUtils.toString(response.getEntity()));
}
finally {
proxyServer.close();
backend.shutdown();
}
trino = new TrinoContainer("trinodb/trino");
trino.start();

int trinoPort = trino.getMappedPort(8080);
proxyPort = trinoPort + 1;
ProxyServerConfiguration config = new ProxyServerConfiguration();
config.setName("test_cluster");
config.setProxyTo("http://localhost:" + trinoPort);
config.setLocalPort(proxyPort);
proxy = new ProxyServer(config, new ProxyHandler());
proxy.start();
}

@Test
public void testCustomHeader()
throws Exception
@AfterAll
public void cleanup()
{
String mockResponseText = "CUSTOM HEADER TEST";
setProxyServer(mockResponseText);
try {
proxyServer.start();
CloseableHttpClient httpclient = HttpClientBuilder.create().build();
HttpUriRequest httpUriRequest = new HttpGet("http://localhost:" + serverPort);
httpUriRequest.setHeader("HEADER1", "FOO");
httpUriRequest.setHeader("HEADER2", "BAR");

HttpResponse response = httpclient.execute(httpUriRequest);
assertEquals(mockResponseText, EntityUtils.toString(response.getEntity()));
RecordedRequest recordedRequest = backend.takeRequest();
assertEquals("FOO", recordedRequest.getHeader("HEADER1"));
assertEquals("BAR", recordedRequest.getHeader("HEADER2"));
}
finally {
proxyServer.close();
backend.shutdown();
}
trino.close();
proxy.close();
}

@Test
public void testLongHeader()
public void testProxyServer()
throws Exception
{
String mockResponseText = "CUSTOM LONG HEADER TEST";
setProxyServer(mockResponseText);
String mockLongHeaderKey = "HEADER_LONG";
// Mockserver has max 8k for HTTP Header values so test with header value larger than default 4k
int headerLength = 5 * 1024;
String mockLongHeaderValue = "x".repeat(headerLength);
try {
proxyServer.start();
CloseableHttpClient httpclient = HttpClientBuilder.create().build();
HttpUriRequest httpUriRequest = new HttpGet("http://localhost:" + serverPort);
httpUriRequest.setHeader(mockLongHeaderKey, mockLongHeaderValue);

HttpResponse response = httpclient.execute(httpUriRequest);
assertEquals(mockResponseText, EntityUtils.toString(response.getEntity()));
RecordedRequest recordedRequest = backend.takeRequest();
assertEquals(mockLongHeaderValue, recordedRequest.getHeader(mockLongHeaderKey));
}
finally {
proxyServer.close();
backend.shutdown();
try (Connection connection = DriverManager.getConnection("jdbc:trino://localhost:" + proxyPort, "test_user", null);
Statement statement = connection.createStatement()) {
ResultSet result = statement.executeQuery("SELECT 'test'");
result.next();
assertThat(result.getString(1)).isEqualTo("test");
}
}

private ProxyServerConfiguration buildConfig(String backendUrl, int localPort)
{
ProxyServerConfiguration config = new ProxyServerConfiguration();
config.setName("MockBackend");
config.setPrefix("/");
config.setPreserveHost("true");
config.setProxyTo(backendUrl);
config.setLocalPort(localPort);
config.setOutputBufferSize(32 * 1024);
config.setResponseHeaderSize(8 * 1024);
config.setRequestHeaderSize(8 * 1024);
config.setRequestBufferSize(16 * 1024); // default 4 * 1024
config.setResponseBufferSize(16 * 1024);
return config;
}

private void setProxyServer(String mockResponseText)
throws IOException
{
int backendPort = 30000 + new Random().nextInt(1000);

backend = new MockWebServer();
backend.enqueue(new MockResponse().setBody(mockResponseText));
backend.play(backendPort);

serverPort = backendPort + 1;
ProxyServerConfiguration config = buildConfig(backend.getUrl("/").toString(), serverPort);
proxyServer = new ProxyServer(config, new ProxyHandler());
}
}
7 changes: 0 additions & 7 deletions pom.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,6 @@
<dep.jakarta.mail.version>2.0.0</dep.jakarta.mail.version>
<dep.jeasy.version>4.1.0</dep.jeasy.version>
<dep.jetty.version>11.0.15</dep.jetty.version>
<dep.mockwebserver.version>1.2.1</dep.mockwebserver.version>
<dep.mockito.version>5.8.0</dep.mockito.version>
<dep.mysqlconnector.version>8.0.17</dep.mysqlconnector.version>
<dep.okhttp.version>3.9.0</dep.okhttp.version>
Expand Down Expand Up @@ -90,12 +89,6 @@
<scope>import</scope>
</dependency>

<dependency>
<groupId>com.squareup.okhttp</groupId>
<artifactId>mockwebserver</artifactId>
<version>${dep.mockwebserver.version}</version>
</dependency>

<dependency>
<groupId>org.apache.commons</groupId>
<artifactId>commons-lang3</artifactId>
Expand Down

0 comments on commit f366966

Please sign in to comment.