Change testAuth and testAuth2 to allocate only one server

Bit-Quill · Nov 14, 2023 · 199002e · 199002e
1 parent 0a93a6a
commit 199002e
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 22 deletions.
diff --git a/java/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/TestBasicAuth.java b/java/flight/flight-core/src/test/java/org/apache/arrow/flight/auth/TestBasicAuth.java
@@ -41,8 +41,10 @@
 import org.apache.arrow.vector.types.Types;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.Schema;
+import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.RepeatedTest;
 import org.junit.jupiter.api.Test;
@@ -56,8 +58,8 @@ public class TestBasicAuth {
   private static final byte[] VALID_TOKEN = "my_token".getBytes(StandardCharsets.UTF_8);
 
   private FlightClient client;
-  private FlightServer server;
-  private BufferAllocator allocator;
+  private static FlightServer server;
+  private static BufferAllocator allocator;
 
   @Test
   public void validAuth() {
@@ -95,7 +97,12 @@ public void didntAuth() {
   }
 
   @BeforeEach
-  public void setup() throws IOException {
+  public void testSetup() throws IOException {
+    client = FlightClient.builder(allocator, server.getLocation()).build();
+  }
+
+  @BeforeAll
+  public static void setup() throws IOException {
     allocator = new RootAllocator(Long.MAX_VALUE);
     final BasicServerAuthHandler.BasicAuthValidator validator = new BasicServerAuthHandler.BasicAuthValidator() {
 
@@ -147,12 +154,16 @@ public void getStream(CallContext context, Ticket ticket, ServerStreamListener l
             }
           }
         }).authHandler(new BasicServerAuthHandler(validator)).build().start();
-    client = FlightClient.builder(allocator, server.getLocation()).build();
   }
 
   @AfterEach
-  public void shutdown() throws Exception {
-    AutoCloseables.close(client, server);
+  public void tearDown() throws Exception {
+    AutoCloseables.close(client);
+  }
+
+  @AfterAll
+  public static void shutdown() throws Exception {
+    AutoCloseables.close(server);
 
     allocator.getChildAllocators().forEach(BufferAllocator::close);
     AutoCloseables.close(allocator);

diff --git a/java/flight/flight-core/src/test/java/org/apache/arrow/flight/auth2/TestBasicAuth2.java b/java/flight/flight-core/src/test/java/org/apache/arrow/flight/auth2/TestBasicAuth2.java
@@ -41,9 +41,9 @@
 import org.apache.arrow.vector.types.Types;
 import org.apache.arrow.vector.types.pojo.Field;
 import org.apache.arrow.vector.types.pojo.Schema;
-import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.AfterAll;
 import org.junit.jupiter.api.Assertions;
-import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.BeforeAll;
 import org.junit.jupiter.api.RepeatedTest;
 import org.junit.jupiter.api.Test;
 
@@ -57,18 +57,18 @@ public class TestBasicAuth2 {
   private static final String NO_USERNAME = "";
   private static final String PASSWORD_1 = "woohoo1";
   private static final String PASSWORD_2 = "woohoo2";
-  private BufferAllocator allocator;
-  private FlightServer server;
-  private FlightClient client;
-  private FlightClient client2;
+  private static BufferAllocator allocator;
+  private static FlightServer server;
+  private static FlightClient client;
+  private static FlightClient client2;
 
-  @BeforeEach
-  public void setup() throws Exception {
+  @BeforeAll
+  public static void setup() throws Exception {
     allocator = new RootAllocator(Long.MAX_VALUE);
     startServerAndClient();
   }
 
-  private FlightProducer getFlightProducer() {
+  private static FlightProducer getFlightProducer() {
     return new NoOpFlightProducer() {
       @Override
       public void listFlights(CallContext context, Criteria criteria,
@@ -99,19 +99,19 @@ public void getStream(CallContext context, Ticket ticket, ServerStreamListener l
     };
   }
 
-  private void startServerAndClient() throws IOException {
+  private static void startServerAndClient() throws IOException {
     final FlightProducer flightProducer = getFlightProducer();
-    this.server = FlightServer
+    server = FlightServer
         .builder(allocator, forGrpcInsecure(LOCALHOST, 0), flightProducer)
         .headerAuthenticator(new GeneratedBearerTokenAuthenticator(
-            new BasicCallHeaderAuthenticator(this::validate)))
+            new BasicCallHeaderAuthenticator(TestBasicAuth2::validate)))
         .build().start();
-    this.client = FlightClient.builder(allocator, server.getLocation())
+    client = FlightClient.builder(allocator, server.getLocation())
         .build();
   }
 
-  @AfterEach
-  public void shutdown() throws Exception {
+  @AfterAll
+  public static void shutdown() throws Exception {
     AutoCloseables.close(client, client2, server);
     client = null;
     client2 = null;
@@ -127,7 +127,7 @@ private void startClient2() throws IOException {
         .build();
   }
 
-  private CallHeaderAuthenticator.AuthResult validate(String username, String password) {
+  private static CallHeaderAuthenticator.AuthResult validate(String username, String password) {
     if (Strings.isNullOrEmpty(username)) {
       throw CallStatus.UNAUTHENTICATED.withDescription("Credentials not supplied.").toRuntimeException();
     }