From a9d3ded8193a872b18826da41ca1104886b6fd40 Mon Sep 17 00:00:00 2001
From: Maximilian Michels <mxm@apache.org>
Date: Tue, 18 Dec 2018 18:53:36 +0100
Subject: [PATCH] Fix port problem / error reporting

---
 .../flink/FlinkJobServerDriverTest.java       | 25 ++++++-------------
 1 file changed, 7 insertions(+), 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
index 11f617005034a..23189c7d3dbc0 100644
--- 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();
-    }
-  }
 }