[SPARK-501] Support CNI network labels: added integration test (apache#167)

Author: susanxhuynh

tests/test_recovery.py

@@ -35,10 +35,7 @@ def test_disconnect_from_master():
                      "--conf", "spark.cores.max=1"])
 
     # Wait until executor is running
-    LOGGER.info("Waiting for executor task to be RUNNING...")
-    shakedown.wait_for(lambda: utils.is_service_ready(LONG_RUNNING_FW_NAME, LONG_RUNNING_FW_NUM_TASKS),
-                       ignore_exceptions=False,
-                       timeout_seconds=600)
+    utils.wait_for_executors_running(LONG_RUNNING_FW_NAME, LONG_RUNNING_FW_NUM_TASKS)
 
     # Block the driver's connection to Mesos master
     framework_info = shakedown.get_service(LONG_RUNNING_FW_NAME)

tests/test_spark.py

@@ -20,6 +20,8 @@
 
 LOGGER = logging.getLogger(__name__)
 THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+SPARK_PI_FW_NAME = "Spark Pi"
+CNI_TEST_NUM_EXECUTORS = 1
 
 
 def setup_module(module):
@@ -113,6 +115,49 @@ def test_cni():
                 "--class", "org.apache.spark.examples.SparkPi"])
 
 
+@pytest.mark.skip("Enable when SPARK-21694 is merged and released in DC/OS Spark")
+@pytest.mark.sanity
+def test_cni_labels():
+    SPARK_EXAMPLES="http://downloads.mesosphere.com/spark/assets/spark-examples_2.11-2.0.1.jar"
+    driver_task_id = utils.submit_job(SPARK_EXAMPLES,
+                               "3000",   # Long enough to examine the Driver's & Executor's task infos
+                               ["--conf", "spark.mesos.network.name=dcos",
+                                "--conf", "spark.mesos.network.labels=key1:val1,key2:val2",
+                                "--conf", "spark.cores.max={}".format(CNI_TEST_NUM_EXECUTORS),
+                                "--class", "org.apache.spark.examples.SparkPi"])
+
+    # Wait until executors are running
+    utils.wait_for_executors_running(SPARK_PI_FW_NAME, CNI_TEST_NUM_EXECUTORS)
+
+    # Check for network name / labels in Driver task info
+    driver_task = shakedown.get_task(driver_task_id, completed=False)
+    _check_task_network_info(driver_task)
+
+    # Check for network name / labels in Executor task info
+    executor_task = shakedown.get_service_tasks(SPARK_PI_FW_NAME)[0]
+    _check_task_network_info(executor_task)
+
+    # Check job output
+    utils.check_job_output(driver_task_id, "Pi is roughly 3")
+
+
+def _check_task_network_info(task):
+    # Expected: "network_infos":[{
+    #   "name":"dcos",
+    #   "labels":{
+    #       "labels":[
+    #           {"key":"key1","value":"val1"},
+    #           {"key":"key2","value":"val2"}]}}]
+    network_info = task['container']['network_infos'][0]
+    assert network_info['name'] == "dcos"
+    labels = network_info['labels']['labels']
+    assert len(labels) == 2
+    assert labels[0]['key'] == "key1"
+    assert labels[0]['value'] == "val1"
+    assert labels[1]['key'] == "key2"
+    assert labels[1]['value'] == "val2"
+
+
 @pytest.mark.sanity
 def test_s3():
     linecount_path = os.path.join(THIS_DIR, 'resources', 'linecount.txt')

tests/utils.py

@@ -206,6 +206,13 @@ def submit_job(app_url, app_args, args=[]):
     return match.group(1)
 
 
+def wait_for_executors_running(framework_name, num_executors):
+    LOGGER.info("Waiting for executor task to be RUNNING...")
+    shakedown.wait_for(lambda: is_service_ready(framework_name, num_executors),
+                       ignore_exceptions=False,
+                       timeout_seconds=600)
+
+
 def _task_log(task_id, filename=None):
     cmd = "dcos task log --completed --lines=1000 {}".format(task_id) + \
           ("" if filename is None else " {}".format(filename))