Fix port problem / error reporting

apache · Dec 18, 2018 · a9d3ded · a9d3ded
1 parent 1c2b08d
commit a9d3ded
Showing 1 changed file with 7 additions and 18 deletions.
diff --git a/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkJobServerDriverTest.java b/runners/flink/src/test/java/org/apache/beam/runners/flink/FlinkJobServerDriverTest.java
@@ -24,9 +24,7 @@
 
 import com.google.common.base.Charsets;
 import java.io.ByteArrayOutputStream;
-import java.io.IOException;
 import java.io.PrintStream;
-import java.net.ServerSocket;
 import org.junit.Test;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -89,28 +87,25 @@ public void testJobServerDriver() throws Exception {
     PrintStream newOut = new PrintStream(baos);
     try {
       System.setErr(newOut);
-      int freePort = getFreePort();
-      int freePort2 = getFreePort();
-      driver =
-          FlinkJobServerDriver.fromParams(
-              new String[] {
-                "--job-port", String.valueOf(freePort),
-                "--artifact-port", String.valueOf(freePort2)
-              });
+      driver = FlinkJobServerDriver.fromParams(new String[] {"--job-port=0", "--artifact-port=0"});
       driverThread = new Thread(driver);
       driverThread.start();
       boolean success = false;
       while (!success) {
         newOut.flush();
         String output = baos.toString(Charsets.UTF_8.name());
-        if (output.contains("JobService started on localhost:" + freePort)
-            && output.contains("ArtifactStagingService started on localhost:" + freePort2)) {
+        if (output.contains("JobService started on localhost:")
+            && output.contains("ArtifactStagingService started on localhost:")) {
           success = true;
         } else {
           Thread.sleep(100);
         }
       }
       assertThat(driverThread.isAlive(), is(true));
+    } catch (Throwable t) {
+      // restore to print exception
+      System.setErr(oldOut);
+      throw t;
     } finally {
       System.setErr(oldOut);
       if (driver != null) {
@@ -122,10 +117,4 @@ public void testJobServerDriver() throws Exception {
       }
     }
   }
-
-  private static int getFreePort() throws IOException {
-    try (ServerSocket socket = new ServerSocket(0)) {
-      return socket.getLocalPort();
-    }
-  }
 }