Skip to content

Commit

Permalink
Improve test coverage for ThriftHttpMetastoreClient
Browse files Browse the repository at this point in the history
  • Loading branch information
@piotrrzysko
piotrrzysko committed Jan 14, 2025
1 parent f548262 commit bd1cda2
Show file tree
Hide file tree
Showing 2 changed files with 58 additions and 30 deletions.
57 changes: 37 additions & 20 deletions ...ve/src/test/java/io/trino/plugin/hive/metastore/thrift/TestThriftHttpMetastoreClient.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,9 @@
import io.opentelemetry.api.OpenTelemetry;
import io.trino.hive.thrift.metastore.Database;
import io.trino.hive.thrift.metastore.NoSuchObjectException;
import io.trino.hive.thrift.metastore.TableMeta;
import io.trino.plugin.base.util.AutoCloseableCloser;
import io.trino.plugin.hive.metastore.thrift.TestingThriftHttpMetastoreServer.TestingThriftRequestsHandler;
import io.trino.testing.TestingNodeManager;
import jakarta.servlet.http.HttpServletRequest;
import org.apache.http.HttpHeaders;
Expand All @@ -34,22 +36,19 @@
import java.util.Optional;
import java.util.function.Consumer;

import static io.trino.plugin.hive.TestingThriftHiveMetastoreBuilder.testingThriftHiveMetastoreBuilder;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;

public class TestThriftHttpMetastoreClient
{
private static final AutoCloseableCloser closer = AutoCloseableCloser.create();
private static ThriftMetastore delegate;

@BeforeAll
public static void setup()
throws Exception
{
File tempDir = Files.createTempDirectory(null).toFile();
tempDir.deleteOnExit();
delegate = testingThriftHiveMetastoreBuilder().metastoreClient(createFakeMetastoreClient()).build(closer::register);
}

@AfterAll
Expand All @@ -59,43 +58,61 @@ public static void tearDown()
closer.close();
}

private static ThriftMetastoreClient createFakeMetastoreClient()
@Test
public void testHttpThriftConnection()
throws Exception
{
return new MockThriftMetastoreClient()
ThriftHttpMetastoreConfig config = new ThriftHttpMetastoreConfig();
config.setAuthenticationMode(ThriftHttpMetastoreConfig.AuthenticationMode.BEARER);
config.setAdditionalHeaders("key1:value1, key2:value2");

TestingThriftRequestsHandler handler = new TestingThriftRequestsHandler()
{
@Override
public Database getDatabase(String databaseName)
public List<String> getAllDatabases()
{
return ImmutableList.of("testDbName");
}

@Override
public Database getDatabase(String name)
throws NoSuchObjectException
{
if (databaseName.equals("testDbName")) {
return new Database(databaseName, "testOwner", "testLocation", Map.of("key", "value"));
if (name.equals("testDbName")) {
return new Database(name, "testOwner", "testLocation", Map.of("key", "value"));
}
throw new NoSuchObjectException("Database does not exist");
}

@Override
public List<String> getAllDatabases()
public List<String> getTables(String databaseName, String pattern)
{
return ImmutableList.of("testDbName");
if (databaseName.equals("testDbName")) {
return ImmutableList.of("testTable1", "testTable2");
}
return ImmutableList.of();
}
};
}

@Test
public void testHttpThriftConnection()
throws Exception
{
ThriftHttpMetastoreConfig config = new ThriftHttpMetastoreConfig();
config.setAuthenticationMode(ThriftHttpMetastoreConfig.AuthenticationMode.BEARER);
config.setAdditionalHeaders("key1:value1, key2:value2");
@Override
public List<String> getTablesByType(String databaseName, String pattern, String tableType)
{
if (databaseName.equals("testDbName")) {
return ImmutableList.of("testTable3", "testTable4");
}
return ImmutableList.of();
}
};

try (TestingThriftHttpMetastoreServer metastoreServer = new TestingThriftHttpMetastoreServer(delegate, new TestRequestHeaderInterceptor())) {
try (TestingThriftHttpMetastoreServer metastoreServer = new TestingThriftHttpMetastoreServer(handler, new TestRequestHeaderInterceptor())) {
ThriftMetastoreClientFactory factory = new HttpThriftMetastoreClientFactory(config, new TestingNodeManager(), OpenTelemetry.noop());
URI metastoreUri = URI.create("http://localhost:" + metastoreServer.getPort());
ThriftMetastoreClient client = factory.create(
metastoreUri, Optional.empty());
assertThat(client.getAllDatabases()).containsExactly("testDbName");
assertThat(client.getDatabase("testDbName")).isEqualTo(new Database("testDbName", "testOwner", "testLocation", Map.of("key", "value")));
assertThat(client.getTableMeta("testDbName"))
.extracting(TableMeta::getTableName)
.containsExactlyInAnyOrder("testTable1", "testTable2", "testTable3", "testTable4");
// negative case
assertThatThrownBy(() -> client.getDatabase("does-not-exist")).isInstanceOf(NoSuchObjectException.class);
}
Expand Down
31 changes: 21 additions & 10 deletions ...src/test/java/io/trino/plugin/hive/metastore/thrift/TestingThriftHttpMetastoreServer.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@
import java.io.Closeable;
import java.io.IOException;
import java.net.URI;
import java.util.Optional;
import java.util.List;
import java.util.function.Consumer;

import static com.google.common.reflect.Reflection.newProxy;
Expand All @@ -46,9 +46,9 @@ public class TestingThriftHttpMetastoreServer
private final LifeCycleManager lifeCycleManager;
private final URI baseUri;

public TestingThriftHttpMetastoreServer(ThriftMetastore delegate, Consumer<HttpServletRequest> requestInterceptor)
public TestingThriftHttpMetastoreServer(TestingThriftRequestsHandler handler, Consumer<HttpServletRequest> requestInterceptor)
{
ThriftHiveMetastore.Iface mockThriftHandler = proxyHandler(delegate, ThriftHiveMetastore.Iface.class);
ThriftHiveMetastore.Iface mockThriftHandler = proxyHandler(handler);
TProcessor processor = new ThriftHiveMetastore.Processor<>(mockThriftHandler);
thriftHttpServlet = new TestingThriftHttpServlet(processor, new TBinaryProtocol.Factory(), requestInterceptor);
Bootstrap app = new Bootstrap(
Expand All @@ -67,14 +67,13 @@ public TestingThriftHttpMetastoreServer(ThriftMetastore delegate, Consumer<HttpS
baseUri = httpServerInfo.getHttpUri();
}

private static <T> T proxyHandler(ThriftMetastore delegate, Class<T> iface)
private static ThriftHiveMetastore.Iface proxyHandler(TestingThriftRequestsHandler handler)
{
return newProxy(iface, (proxy, method, args) -> switch (method.getName()) {
case "getAllDatabases" -> delegate.getAllDatabases();
case "getDatabase" -> {
Optional<Database> optionalDatabase = delegate.getDatabase(args[0].toString());
yield optionalDatabase.orElseThrow(() -> new NoSuchObjectException(""));
}
return newProxy(ThriftHiveMetastore.Iface.class, (_, method, args) -> switch (method.getName()) {
case "getAllDatabases" -> handler.getAllDatabases();
case "getDatabase" -> handler.getDatabase(args[0].toString());
case "getTables" -> handler.getTables(args[0].toString(), args[1].toString());
case "getTablesByType" -> handler.getTablesByType(args[0].toString(), args[1].toString(), args[2].toString());
default -> throw new UnsupportedOperationException();
});
}
Expand Down Expand Up @@ -114,4 +113,16 @@ protected void doPost(HttpServletRequest request,
super.doPost(request, response);
}
}

public interface TestingThriftRequestsHandler
{
List<String> getAllDatabases();

Database getDatabase(String name)
throws NoSuchObjectException;

List<String> getTables(String databaseName, String pattern);

List<String> getTablesByType(String databaseName, String pattern, String tableType);
}
}

0 comments on commit bd1cda2

Please sign in to comment.