diff --git a/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/FlightSqlStatelessExample.java b/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/FlightSqlStatelessExample.java index 8cd5d622f7381..02f679b3efaf4 100644 --- a/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/FlightSqlStatelessExample.java +++ b/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/example/FlightSqlStatelessExample.java @@ -42,14 +42,11 @@ import org.apache.arrow.flight.CallStatus; import org.apache.arrow.flight.FlightDescriptor; import org.apache.arrow.flight.FlightInfo; -import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightStream; import org.apache.arrow.flight.Location; import org.apache.arrow.flight.PutResult; import org.apache.arrow.flight.sql.FlightSqlProducer; import org.apache.arrow.memory.ArrowBuf; -import org.apache.arrow.memory.BufferAllocator; -import org.apache.arrow.memory.RootAllocator; import org.apache.arrow.vector.VectorLoader; import org.apache.arrow.vector.VectorSchemaRoot; import org.apache.arrow.vector.VectorUnloader; @@ -70,17 +67,6 @@ public class FlightSqlStatelessExample extends FlightSqlExample { private static final Logger LOGGER = getLogger(FlightSqlStatelessExample.class); - public static void main(String[] args) throws Exception { - Location location = Location.forGrpcInsecure("localhost", 55555); - final FlightSqlStatelessExample example = new FlightSqlStatelessExample(location); - Location listenLocation = Location.forGrpcInsecure("0.0.0.0", 55555); - try (final BufferAllocator allocator = new RootAllocator(); - final FlightServer server = FlightServer.builder(allocator, listenLocation, example).build()) { - server.start(); - server.awaitTermination(); - } - } - public FlightSqlStatelessExample(final Location location) { super(location); } diff --git a/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java b/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java index 264caf079658f..b99a17221fc6a 100644 --- a/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java +++ b/java/flight/flight-sql/src/test/java/org/apache/arrow/flight/sql/test/TestFlightSqlStateless.java @@ -22,6 +22,7 @@ import static org.hamcrest.CoreMatchers.*; import org.apache.arrow.flight.FlightClient; +import org.apache.arrow.flight.FlightEndpoint; import org.apache.arrow.flight.FlightInfo; import org.apache.arrow.flight.FlightServer; import org.apache.arrow.flight.FlightStream; @@ -81,15 +82,17 @@ public void testSimplePreparedStatementResultsWithParameterBinding() throws Exce prepare.setParameters(insertRoot); final FlightInfo flightInfo = prepare.execute(); - try (FlightStream stream = sqlClient.getStream(flightInfo.getEndpoints().get(0).getTicket())) { - // TODO: root is null and getSchema hangs when run as complete suite. - // This works when run as an individual test. - final VectorSchemaRoot root = stream.getRoot(); - final Schema schema = root.getSchema(); - Assertions.assertAll( - () -> MatcherAssert.assertThat(schema, is(SCHEMA_INT_TABLE)), - () -> MatcherAssert.assertThat(getResults(stream), is(EXPECTED_RESULTS_FOR_PARAMETER_BINDING)) - ); + for (FlightEndpoint endpoint: flightInfo.getEndpoints()) { + try (FlightStream stream = sqlClient.getStream(endpoint.getTicket())) { + // TODO: root is null and getSchema hangs when run as complete suite. + // This works when run as an individual test. + final VectorSchemaRoot root = stream.getRoot(); + final Schema schema = root.getSchema(); + Assertions.assertAll( + () -> MatcherAssert.assertThat(schema, is(SCHEMA_INT_TABLE)), + () -> MatcherAssert.assertThat(getResults(stream), is(EXPECTED_RESULTS_FOR_PARAMETER_BINDING)) + ); + } } } }