]>
Commit | Line | Data |
---|---|---|
1d09f67e TL |
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 | } |