[SPARK-14391][LAUNCHER] Fix launcher communication test, take 2.

There's actually a race here: the state of the handler was changed before
the connection was set, so the test code could be notified of the state
change, wake up, and still see the connection as null, triggering the assert.

Author: Marcelo Vanzin <vanzin@cloudera.com>

Closes #12785 from vanzin/SPARK-14391.

2 files changed, 2 insertions, 3 deletions

diff --git a/launcher/src/main/java/org/apache/spark/launcher/LauncherServer.java b/launcher/src/main/java/org/apache/spark/launcher/LauncherServer.java
index 69fbf4387b..e3413fd665 100644
--- a/launcher/src/main/java/org/apache/spark/launcher/LauncherServer.java
+++ b/launcher/src/main/java/org/apache/spark/launcher/LauncherServer.java
@@ -298,8 +298,8 @@ class LauncherServer implements Closeable {
           Hello hello = (Hello) msg;
           ChildProcAppHandle handle = pending.remove(hello.secret);
           if (handle != null) {
-            handle.setState(SparkAppHandle.State.CONNECTED);
             handle.setConnection(this);
+            handle.setState(SparkAppHandle.State.CONNECTED);
             this.handle = handle;
           } else {
             throw new IllegalArgumentException("Received Hello for unknown client.");
diff --git a/launcher/src/test/java/org/apache/spark/launcher/LauncherServerSuite.java b/launcher/src/test/java/org/apache/spark/launcher/LauncherServerSuite.java
index 22b834c575..bfe1fcc87f 100644
--- a/launcher/src/test/java/org/apache/spark/launcher/LauncherServerSuite.java
+++ b/launcher/src/test/java/org/apache/spark/launcher/LauncherServerSuite.java
@@ -26,7 +26,6 @@ import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 
-import org.junit.Ignore;
 import org.junit.Test;
 import static org.junit.Assert.*;
 
@@ -62,7 +61,7 @@ public class LauncherServerSuite extends BaseSuite {
     }
   }
 
-  @Ignore
+  @Test
   public void testCommunication() throws Exception {
     ChildProcAppHandle handle = LauncherServer.newAppHandle();
     TestClient client = null;