ceph/src/arrow/java/flight/flight-core/src/test/java/org/apache/arrow/flight/FlightTestUtil.java

   1 /*
   2  * Licensed to the Apache Software Foundation (ASF) under one or more
   3  * contributor license agreements.  See the NOTICE file distributed with
   4  * this work for additional information regarding copyright ownership.
   5  * The ASF licenses this file to You under the Apache License, Version 2.0
   6  * (the "License"); you may not use this file except in compliance with
   7  * the License.  You may obtain a copy of the License at
   8  *
   9  *    http://www.apache.org/licenses/LICENSE-2.0
  10  *
  11  * Unless required by applicable law or agreed to in writing, software
  12  * distributed under the License is distributed on an "AS IS" BASIS,
  13  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  14  * See the License for the specific language governing permissions and
  15  * limitations under the License.
  16  */
  17
  18 package org.apache.arrow.flight;
  19
  20 import java.io.File;
  21 import java.io.IOException;
  22 import java.lang.reflect.InvocationTargetException;
  23 import java.nio.file.Path;
  24 import java.nio.file.Paths;
  25 import java.util.Arrays;
  26 import java.util.List;
  27 import java.util.Objects;
  28 import java.util.Random;
  29 import java.util.function.Function;
  30
  31 import org.junit.Assert;
  32 import org.junit.jupiter.api.Assertions;
  33 import org.junit.jupiter.api.function.Executable;
  34
  35 /**
  36  * Utility methods and constants for testing flight servers.
  37  */
  38 public class FlightTestUtil {
  39
  40   private static final Random RANDOM = new Random();
  41
  42   public static final String LOCALHOST = "localhost";
  43   public static final String TEST_DATA_ENV_VAR = "ARROW_TEST_DATA";
  44   public static final String TEST_DATA_PROPERTY = "arrow.test.dataRoot";
  45
  46   /**
  47    * Returns a a FlightServer (actually anything that is startable)
  48    * that has been started bound to a random port.
  49    */
  50   public static <T> T getStartedServer(Function<Location, T> newServerFromLocation) throws IOException {
  51     IOException lastThrown = null;
  52     T server = null;
  53     for (int x = 0; x < 3; x++) {
  54       final int port = 49152 + RANDOM.nextInt(5000);
  55       final Location location = Location.forGrpcInsecure(LOCALHOST, port);
  56       lastThrown = null;
  57       try {
  58         server = newServerFromLocation.apply(location);
  59         try {
  60           server.getClass().getMethod("start").invoke(server);
  61         } catch (NoSuchMethodException | IllegalAccessException e) {
  62           throw new IllegalArgumentException("Couldn't call start method on object.", e);
  63         }
  64         break;
  65       } catch (InvocationTargetException e) {
  66         if (e.getTargetException() instanceof IOException) {
  67           lastThrown = (IOException) e.getTargetException();
  68         } else {
  69           throw (RuntimeException) e.getTargetException();
  70         }
  71       }
  72     }
  73     if (lastThrown != null) {
  74       throw lastThrown;
  75     }
  76     return server;
  77   }
  78
  79   static Path getTestDataRoot() {
  80     String path = System.getenv(TEST_DATA_ENV_VAR);
  81     if (path == null) {
  82       path = System.getProperty(TEST_DATA_PROPERTY);
  83     }
  84     return Paths.get(Objects.requireNonNull(path,
  85         String.format("Could not find test data path. Set the environment variable %s or the JVM property %s.",
  86             TEST_DATA_ENV_VAR, TEST_DATA_PROPERTY)));
  87   }
  88
  89   static Path getFlightTestDataRoot() {
  90     return getTestDataRoot().resolve("flight");
  91   }
  92
  93   static Path exampleTlsRootCert() {
  94     return getFlightTestDataRoot().resolve("root-ca.pem");
  95   }
  96
  97   static List<CertKeyPair> exampleTlsCerts() {
  98     final Path root = getFlightTestDataRoot();
  99     return Arrays.asList(new CertKeyPair(root.resolve("cert0.pem").toFile(), root.resolve("cert0.pkcs1").toFile()),
 100         new CertKeyPair(root.resolve("cert1.pem").toFile(), root.resolve("cert1.pkcs1").toFile()));
 101   }
 102
 103   static boolean isEpollAvailable() {
 104     try {
 105       Class<?> epoll = Class.forName("io.netty.channel.epoll.Epoll");
 106       return (Boolean) epoll.getMethod("isAvailable").invoke(null);
 107     } catch (ClassNotFoundException | NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
 108       return false;
 109     }
 110   }
 111
 112   static boolean isKqueueAvailable() {
 113     try {
 114       Class<?> kqueue = Class.forName("io.netty.channel.kqueue.KQueue");
 115       return (Boolean) kqueue.getMethod("isAvailable").invoke(null);
 116     } catch (ClassNotFoundException | NoSuchMethodException | IllegalAccessException | InvocationTargetException e) {
 117       return false;
 118     }
 119   }
 120
 121   static boolean isNativeTransportAvailable() {
 122     return isEpollAvailable() || isKqueueAvailable();
 123   }
 124
 125   /**
 126    * Assert that the given runnable fails with a Flight exception of the given code.
 127    * @param code The expected Flight status code.
 128    * @param r The code to run.
 129    * @return The thrown status.
 130    */
 131   public static CallStatus assertCode(FlightStatusCode code, Executable r) {
 132     final FlightRuntimeException ex = Assertions.assertThrows(FlightRuntimeException.class, r);
 133     Assert.assertEquals(code, ex.status().code());
 134     return ex.status();
 135   }
 136
 137   public static class CertKeyPair {
 138
 139     public final File cert;
 140     public final File key;
 141
 142     public CertKeyPair(File cert, File key) {
 143       this.cert = cert;
 144       this.key = key;
 145     }
 146   }
 147
 148   private FlightTestUtil() {
 149   }
 150 }