Added a few more test cases (#11)

* Added a few more test cases

* Fixed name

* Combined the tests into one

* Added a test case for custom driver pod name

* Combined the tests into one

Author: liyinan926

integration-test/src/test/scala/org/apache/spark/deploy/k8s/integrationtest/KubernetesSuite.scala

@@ -22,12 +22,12 @@ import java.util.UUID
 import java.util.regex.Pattern
 
 import com.google.common.io.PatternFilenameFilter
+import io.fabric8.kubernetes.api.model.Pod
 import org.scalatest.{BeforeAndAfter, BeforeAndAfterAll, FunSuite}
 import org.scalatest.concurrent.{Eventually, PatienceConfiguration}
 import org.scalatest.time.{Minutes, Seconds, Span}
 
 import org.apache.spark.deploy.k8s.integrationtest.backend.IntegrationTestBackendFactory
-import org.apache.spark.deploy.k8s.integrationtest.constants.MINIKUBE_TEST_BACKEND
 import org.apache.spark.deploy.k8s.integrationtest.constants.SPARK_DISTRO_PATH
 
 private[spark] class KubernetesSuite extends FunSuite with BeforeAndAfterAll with BeforeAndAfter {
@@ -66,20 +66,56 @@ private[spark] class KubernetesSuite extends FunSuite with BeforeAndAfterAll wit
     runSparkPiAndVerifyCompletion()
   }
 
+  test("Run SparkPi with a master URL without a scheme.") {
+    val url = kubernetesTestComponents.kubernetesClient.getMasterUrl
+    sparkAppConf.set("spark.master", s"k8s://${url.getHost}:${url.getPort}")
+    runSparkPiAndVerifyCompletion()
+  }
+
+  test("Run SparkPi with custom driver pod name, labels, annotations, and environment variables.") {
+    sparkAppConf
+      .set("spark.kubernetes.driver.pod.name", "spark-integration-spark-pi")
+      .set("spark.kubernetes.driver.label.label1", "label1-value")
+      .set("spark.kubernetes.driver.label.label2", "label2-value")
+      .set("spark.kubernetes.driver.annotation.annotation1", "annotation1-value")
+      .set("spark.kubernetes.driver.annotation.annotation2", "annotation2-value")
+      .set("spark.kubernetes.driverEnv.ENV1", "VALUE1")
+      .set("spark.kubernetes.driverEnv.ENV2", "VALUE2")
+    runSparkPiAndVerifyCompletion(driverPodChecker = (driverPod: Pod) => {
+      doBasicDriverPodCheck(driverPod)
+      assert(driverPod.getMetadata.getName === "spark-integration-spark-pi")
+
+      assert(driverPod.getMetadata.getLabels.get("label1") === "label1-value")
+      assert(driverPod.getMetadata.getLabels.get("label2") === "label2-value")
+      assert(driverPod.getMetadata.getAnnotations.get("annotation1") === "annotation1-value")
+      assert(driverPod.getMetadata.getAnnotations.get("annotation2") === "annotation2-value")
+
+      val driverContainer = driverPod.getSpec.getContainers.get(0)
+      assert(driverContainer.getEnv.size() == 2)
+      assert(driverContainer.getEnv.get(0).getName === "ENV1")
+      assert(driverContainer.getEnv.get(0).getValue === "VALUE1")
+      assert(driverContainer.getEnv.get(1).getName === "ENV2")
+      assert(driverContainer.getEnv.get(1).getValue === "VALUE2")
+    })
+  }
+
   private def runSparkPiAndVerifyCompletion(
-      appResource: String = CONTAINER_LOCAL_SPARK_DISTRO_EXAMPLES_JAR): Unit = {
+      appResource: String = CONTAINER_LOCAL_SPARK_DISTRO_EXAMPLES_JAR,
+      driverPodChecker: Pod => Unit = doBasicDriverPodCheck): Unit = {
     runSparkApplicationAndVerifyCompletion(
-        appResource,
-        SPARK_PI_MAIN_CLASS,
-        Seq("Pi is roughly 3"),
-        Array.empty[String])
+      appResource,
+      SPARK_PI_MAIN_CLASS,
+      Seq("Pi is roughly 3"),
+      Array.empty[String],
+      driverPodChecker)
   }
 
   private def runSparkApplicationAndVerifyCompletion(
       appResource: String,
       mainClass: String,
       expectedLogOnCompletion: Seq[String],
-      appArgs: Array[String]): Unit = {
+      appArgs: Array[String],
+      driverPodChecker: Pod => Unit): Unit = {
     val appArguments = SparkAppArguments(
       mainAppResource = appResource,
       mainClass = mainClass)
@@ -90,6 +126,7 @@ private[spark] class KubernetesSuite extends FunSuite with BeforeAndAfterAll wit
       .list()
       .getItems
       .get(0)
+    driverPodChecker(driverPod)
     Eventually.eventually(TIMEOUT, INTERVAL) {
       expectedLogOnCompletion.foreach { e =>
         assert(kubernetesTestComponents.kubernetesClient
@@ -100,6 +137,10 @@ private[spark] class KubernetesSuite extends FunSuite with BeforeAndAfterAll wit
       }
     }
   }
+
+  private def doBasicDriverPodCheck(driverPod: Pod): Unit = {
+    assert(driverPod.getMetadata.getLabels.get("spark-role") === "driver")
+  }
 }
 
 private[spark] object KubernetesSuite {