From 96027c71c8fb451af8f1fb3ab6a2c882f367099a Mon Sep 17 00:00:00 2001
From: Patrick Wendell <pwendell@gmail.com>
Date: Mon, 28 Apr 2014 11:54:49 -0700
Subject: [PATCH] Test fixes

---
 .../src/main/scala/org/apache/spark/deploy/SparkSubmit.scala | 5 +++--
 .../scala/org/apache/spark/deploy/SparkSubmitSuite.scala     | 5 +++--
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala b/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala
index 1b88f3343177d..85a96efdbb17d 100644
--- a/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala
+++ b/core/src/main/scala/org/apache/spark/deploy/SparkSubmit.scala
@@ -124,12 +124,13 @@ object SparkSubmit {
       }
     } else if (clusterManager == YARN) {
       childMainClass = "org.apache.spark.deploy.yarn.Client"
-      // Special flag to avoid deprecation warnings at the client
-      sysProps("SPARK_SUBMIT") = "true"
       childArgs += ("--jar", appArgs.primaryResource)
       childArgs += ("--class", appArgs.mainClass)
     }
 
+    // Special flag to avoid deprecation warnings at the client
+    sysProps("SPARK_SUBMIT") = "true"
+
     val options = List[OptionAssigner](
       new OptionAssigner(appArgs.master, ALL_CLUSTER_MGRS, false, sysProp = "spark.master"),
       new OptionAssigner(appArgs.driverExtraClassPath, STANDALONE | YARN, true,
diff --git a/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala b/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
index 10a65c75cc621..b3541b4a40b79 100644
--- a/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
+++ b/core/src/test/scala/org/apache/spark/deploy/SparkSubmitSuite.scala
@@ -122,7 +122,7 @@ class SparkSubmitSuite extends FunSuite with ShouldMatchers {
     childArgsStr should include ("--num-executors 6")
     mainClass should be ("org.apache.spark.deploy.yarn.Client")
     classpath should have length (0)
-    sysProps should have size (0)
+    sysProps should have size (1)
   }
 
   test("handles YARN client mode") {
@@ -146,6 +146,7 @@ class SparkSubmitSuite extends FunSuite with ShouldMatchers {
     sysProps("spark.yarn.dist.files") should be ("file1.txt,file2.txt")
     sysProps("spark.yarn.dist.archives") should be ("archive1.txt,archive2.txt")
     sysProps("spark.executor.instances") should be ("6")
+    sysProps("SPARK_SUBMIT") should be ("true")
   }
 
   test("handles standalone cluster mode") {
@@ -159,7 +160,7 @@ class SparkSubmitSuite extends FunSuite with ShouldMatchers {
     childArgsStr should include ("launch spark://h:p thejar.jar org.SomeClass arg1 arg2")
     mainClass should be ("org.apache.spark.deploy.Client")
     classpath should have length (0)
-    sysProps should have size (1) // contains --jar entry
+    sysProps should have size (2) // contains --jar entry and SPARK_SUBMIT
   }
 
   test("handles standalone client mode") {