mock the insecure gateway, clean up everything related to creating insecure gateways

Author: squito

core/src/main/scala/org/apache/spark/api/python/PythonGatewayServer.scala

@@ -43,17 +43,12 @@ private[spark] object PythonGatewayServer extends Logging {
     // with the same secret, in case the app needs callbacks from the JVM to the underlying
     // python processes.
     val localhost = InetAddress.getLoopbackAddress()
-    val builder = new GatewayServer.GatewayServerBuilder()
+    val gatewayServer: GatewayServer = new GatewayServer.GatewayServerBuilder()
+      .authToken(secret)
       .javaPort(0)
       .javaAddress(localhost)
       .callbackClient(GatewayServer.DEFAULT_PYTHON_PORT, localhost, secret)
-    if (sys.env.getOrElse("_PYSPARK_CREATE_INSECURE_GATEWAY", "0") != "1") {
-      builder.authToken(secret)
-    } else {
-      assert(sys.env.getOrElse("SPARK_TESTING", "0") == "1",
-        "Creating insecure Java gateways only allowed for testing")
-    }
-    val gatewayServer: GatewayServer = builder.build()
+      .build()
 
     gatewayServer.start()
     val boundPort: Int = gatewayServer.getListeningPort

python/pyspark/java_gateway.py

@@ -41,20 +41,8 @@ def launch_gateway(conf=None):
     """
     launch jvm gateway
     :param conf: spark configuration passed to spark-submit
-    :return: a JVM gateway
+    :return:
     """
-    return _launch_gateway(conf)
-
-
-def _launch_gateway(conf=None, insecure=False):
-    """
-    launch jvm gateway
-    :param conf: spark configuration passed to spark-submit
-    :param insecure: True to create an insecure gateway; only for testing
-    :return: a JVM gateway
-    """
-    if insecure and os.environ.get("SPARK_TESTING", "0") != "1":
-        raise ValueError("creating insecure gateways is only for testing")
     if "PYSPARK_GATEWAY_PORT" in os.environ:
         gateway_port = int(os.environ["PYSPARK_GATEWAY_PORT"])
         gateway_secret = os.environ["PYSPARK_GATEWAY_SECRET"]
@@ -86,8 +74,6 @@ def _launch_gateway(conf=None, insecure=False):
 
             env = dict(os.environ)
             env["_PYSPARK_DRIVER_CONN_INFO_PATH"] = conn_info_file
-            if insecure:
-                env["_PYSPARK_CREATE_INSECURE_GATEWAY"] = "1"
 
             # Launch the Java gateway.
             # We open a pipe to stdin so that the Java gateway can die when the pipe is broken
@@ -130,10 +116,9 @@ def killChild():
             atexit.register(killChild)
 
     # Connect to the gateway
-    gateway_params = GatewayParameters(port=gateway_port, auto_convert=True)
-    if not insecure:
-        gateway_params.auth_token = gateway_secret
-    gateway = JavaGateway(gateway_parameters=gateway_params)
+    gateway = JavaGateway(
+        gateway_parameters=GatewayParameters(port=gateway_port, auth_token=gateway_secret,
+                                             auto_convert=True))
 
     # Import the classes used by PySpark
     java_import(gateway.jvm, "org.apache.spark.SparkConf")

python/pyspark/tests/test_context.py

@@ -20,9 +20,15 @@
 import threading
 import time
 import unittest
+import sys
+if sys.version >= '3':
+    from unittest.mock import MagicMock
+else:
+    from mock import MagicMock
+
+
 
 from pyspark import SparkFiles, SparkContext
-from pyspark.java_gateway import _launch_gateway
 from pyspark.testing.utils import ReusedPySparkTestCase, PySparkTestCase, QuietTest, SPARK_HOME
 
 
@@ -250,16 +256,12 @@ def test_startTime(self):
     def test_forbid_insecure_gateway(self):
         # Fail immediately if you try to create a SparkContext
         # with an insecure gateway
-        gateway = _launch_gateway(insecure=True)
-        log4j = gateway.jvm.org.apache.log4j
-        old_level = log4j.LogManager.getRootLogger().getLevel()
-        try:
-            log4j.LogManager.getRootLogger().setLevel(log4j.Level.FATAL)
-            with self.assertRaises(Exception) as context:
-                SparkContext(gateway=gateway)
-            self.assertIn("insecure Py4j gateway", str(context.exception))
-        finally:
-            log4j.LogManager.getRootLogger().setLevel(old_level)
+        mock_insecure_gateway = MagicMock()
+        mock_insecure_gateway.gateway_parameters = MagicMock()
+        mock_insecure_gateway.gateway_parameters.auth_token = None
+        with self.assertRaises(Exception) as context:
+            SparkContext(gateway=mock_insecure_gateway)
+        self.assertIn("insecure Py4j gateway", str(context.exception))
 
 
 if __name__ == "__main__":