snowpark-connect 0.24.0__py3-none-any.whl → 0.26.0__py3-none-any.whl

snowflake/snowpark_connect/includes/python/pyspark/tests/test_broadcast.py

@@ -1,196 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one or more
-# contributor license agreements.  See the NOTICE file distributed with
-# this work for additional information regarding copyright ownership.
-# The ASF licenses this file to You under the Apache License, Version 2.0
-# (the "License"); you may not use this file except in compliance with
-# the License.  You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-import os
-import pickle
-import random
-import time
-import tempfile
-import unittest
-
-from py4j.protocol import Py4JJavaError
-
-from pyspark import SparkConf, SparkContext, Broadcast
-from pyspark.java_gateway import launch_gateway
-from pyspark.serializers import ChunkedStream
-from pyspark.sql import SparkSession, Row
-
-
-class BroadcastTest(unittest.TestCase):
-    def tearDown(self):
-        if getattr(self, "sc", None) is not None:
-            self.sc.stop()
-            self.sc = None
-
-    def _test_encryption_helper(self, vs):
-        """
-        Creates a broadcast variables for each value in vs, and runs a simple job to make sure the
-        value is the same when it's read in the executors.  Also makes sure there are no task
-        failures.
-        """
-        bs = [self.sc.broadcast(value=v) for v in vs]
-        exec_values = self.sc.parallelize(range(2)).map(lambda x: [b.value for b in bs]).collect()
-        for ev in exec_values:
-            self.assertEqual(ev, vs)
-        # make sure there are no task failures
-        status = self.sc.statusTracker()
-        for jid in status.getJobIdsForGroup():
-            for sid in status.getJobInfo(jid).stageIds:
-                stage_info = status.getStageInfo(sid)
-                self.assertEqual(0, stage_info.numFailedTasks)
-
-    def _test_multiple_broadcasts(self, *extra_confs):
-        """
-        Test broadcast variables make it OK to the executors.  Tests multiple broadcast variables,
-        and also multiple jobs.
-        """
-        conf = SparkConf()
-        for key, value in extra_confs:
-            conf.set(key, value)
-        conf.setMaster("local-cluster[2,1,1024]")
-        self.sc = SparkContext(conf=conf)
-        self._test_encryption_helper([5])
-        self._test_encryption_helper([5, 10, 20])
-
-    def test_broadcast_with_encryption(self):
-        self._test_multiple_broadcasts(("spark.io.encryption.enabled", "true"))
-
-    def test_broadcast_no_encryption(self):
-        self._test_multiple_broadcasts()
-
-    def _test_broadcast_on_driver(self, *extra_confs):
-        conf = SparkConf()
-        for key, value in extra_confs:
-            conf.set(key, value)
-        conf.setMaster("local-cluster[2,1,1024]")
-        self.sc = SparkContext(conf=conf)
-        bs = self.sc.broadcast(value=5)
-        self.assertEqual(5, bs.value)
-
-    def test_broadcast_value_driver_no_encryption(self):
-        self._test_broadcast_on_driver()
-
-    def test_broadcast_value_driver_encryption(self):
-        self._test_broadcast_on_driver(("spark.io.encryption.enabled", "true"))
-
-    def test_broadcast_value_against_gc(self):
-        # Test broadcast value against gc.
-        conf = SparkConf()
-        conf.setMaster("local[1,1]")
-        conf.set("spark.memory.fraction", "0.0001")
-        self.sc = SparkContext(conf=conf)
-        b = self.sc.broadcast([100])
-        try:
-            res = self.sc.parallelize([0], 1).map(lambda x: 0 if x == 0 else b.value[0]).collect()
-            self.assertEqual([0], res)
-            self.sc._jvm.java.lang.System.gc()
-            time.sleep(5)
-            res = self.sc.parallelize([1], 1).map(lambda x: 0 if x == 0 else b.value[0]).collect()
-            self.assertEqual([100], res)
-        finally:
-            b.destroy()
-
-    def test_broadcast_when_sc_none(self):
-        # SPARK-39029 : Test case to improve test coverage of broadcast.py
-        # It tests the case when SparkContext is none and Broadcast is called at executor
-        conf = SparkConf()
-        conf.setMaster("local-cluster[2,1,1024]")
-        self.sc = SparkContext(conf=conf)
-        bs = self.sc.broadcast([10])
-        bs_sc_none = Broadcast(sc=None, path=bs._path)
-        self.assertEqual(bs_sc_none.value, [10])
-
-    def test_broadcast_for_error_condition(self):
-        # SPARK-39029: Test case to improve test coverage of broadcast.py
-        # It tests the case when broadcast should raise error .
-        conf = SparkConf()
-        conf.setMaster("local-cluster[2,1,1024]")
-        self.sc = SparkContext(conf=conf)
-        bs = self.sc.broadcast([1])
-        with self.assertRaisesRegex(pickle.PickleError, "Could.*not.*serialize.*broadcast"):
-            self.sc.broadcast(self.sc)
-        with self.assertRaisesRegex(Py4JJavaError, "RuntimeError.*Broadcast.*destroyed.*driver"):
-            self.sc.parallelize([1]).map(lambda x: bs.destroy()).collect()
-        with self.assertRaisesRegex(Py4JJavaError, "RuntimeError.*Broadcast.*unpersisted.*driver"):
-            self.sc.parallelize([1]).map(lambda x: bs.unpersist()).collect()
-
-    def test_broadcast_in_udfs_with_encryption(self):
-        conf = SparkConf()
-        conf.set("spark.io.encryption.enabled", "true")
-        conf.setMaster("local-cluster[2,1,1024]")
-        self.sc = SparkContext(conf=conf)
-        bar = {"a": "aa", "b": "bb"}
-        foo = self.sc.broadcast(bar)
-        spark = SparkSession(self.sc)
-        spark.udf.register("MYUDF", lambda x: foo.value[x] if x else "")
-        sel = spark.sql("SELECT MYUDF('a') AS a, MYUDF('b') AS b")
-        self.assertEqual(sel.collect(), [Row(a="aa", b="bb")])
-        spark.stop()
-
-
-class BroadcastFrameProtocolTest(unittest.TestCase):
-    @classmethod
-    def setUpClass(cls):
-        gateway = launch_gateway(SparkConf())
-        cls._jvm = gateway.jvm
-        cls.longMessage = True
-        random.seed(42)
-
-    def _test_chunked_stream(self, data, py_buf_size):
-        # write data using the chunked protocol from python.
-        chunked_file = tempfile.NamedTemporaryFile(delete=False)
-        dechunked_file = tempfile.NamedTemporaryFile(delete=False)
-        dechunked_file.close()
-        try:
-            out = ChunkedStream(chunked_file, py_buf_size)
-            out.write(data)
-            out.close()
-            # now try to read it in java
-            jin = self._jvm.java.io.FileInputStream(chunked_file.name)
-            jout = self._jvm.java.io.FileOutputStream(dechunked_file.name)
-            self._jvm.DechunkedInputStream.dechunkAndCopyToOutput(jin, jout)
-            # java should have decoded it back to the original data
-            self.assertEqual(len(data), os.stat(dechunked_file.name).st_size)
-            with open(dechunked_file.name, "rb") as f:
-                byte = f.read(1)
-                idx = 0
-                while byte:
-                    self.assertEqual(data[idx], bytearray(byte)[0], msg="idx = " + str(idx))
-                    byte = f.read(1)
-                    idx += 1
-        finally:
-            os.unlink(chunked_file.name)
-            os.unlink(dechunked_file.name)
-
-    def test_chunked_stream(self):
-        def random_bytes(n):
-            return bytearray(random.getrandbits(8) for _ in range(n))
-
-        for data_length in [1, 10, 100, 10000]:
-            for buffer_length in [1, 2, 5, 8192]:
-                self._test_chunked_stream(random_bytes(data_length), buffer_length)
-
-
-if __name__ == "__main__":
-    from pyspark.tests.test_broadcast import *  # noqa: F401
-
-    try:
-        import xmlrunner
-
-        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
-    except ImportError:
-        testRunner = None
-    unittest.main(testRunner=testRunner, verbosity=2)

snowflake/snowpark_connect/includes/python/pyspark/tests/test_conf.py

@@ -1,44 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one or more
-# contributor license agreements.  See the NOTICE file distributed with
-# this work for additional information regarding copyright ownership.
-# The ASF licenses this file to You under the Apache License, Version 2.0
-# (the "License"); you may not use this file except in compliance with
-# the License.  You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-import random
-import unittest
-
-from pyspark import SparkContext, SparkConf
-
-
-class ConfTests(unittest.TestCase):
-    def test_memory_conf(self):
-        memoryList = ["1T", "1G", "1M", "1024K"]
-        for memory in memoryList:
-            sc = SparkContext(conf=SparkConf().set("spark.python.worker.memory", memory))
-            lst = list(range(1024))
-            random.shuffle(lst)
-            rdd = sc.parallelize(lst, 4)
-            self.assertEqual(sorted(lst), rdd.sortBy(lambda x: x).collect())
-            sc.stop()
-
-
-if __name__ == "__main__":
-    from pyspark.tests.test_conf import *  # noqa: F401
-
-    try:
-        import xmlrunner
-
-        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
-    except ImportError:
-        testRunner = None
-    unittest.main(testRunner=testRunner, verbosity=2)

snowflake/snowpark_connect/includes/python/pyspark/tests/test_context.py

@@ -1,346 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one or more
-# contributor license agreements.  See the NOTICE file distributed with
-# this work for additional information regarding copyright ownership.
-# The ASF licenses this file to You under the Apache License, Version 2.0
-# (the "License"); you may not use this file except in compliance with
-# the License.  You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-import os
-import shutil
-import stat
-import tempfile
-import threading
-import time
-import unittest
-from collections import namedtuple
-
-from pyspark import SparkConf, SparkFiles, SparkContext
-from pyspark.testing.utils import ReusedPySparkTestCase, PySparkTestCase, QuietTest, SPARK_HOME
-
-
-class CheckpointTests(ReusedPySparkTestCase):
-    def setUp(self):
-        self.checkpointDir = tempfile.NamedTemporaryFile(delete=False)
-        os.unlink(self.checkpointDir.name)
-        self.sc.setCheckpointDir(self.checkpointDir.name)
-
-    def tearDown(self):
-        shutil.rmtree(self.checkpointDir.name)
-
-    def test_basic_checkpointing(self):
-        parCollection = self.sc.parallelize([1, 2, 3, 4])
-        flatMappedRDD = parCollection.flatMap(lambda x: range(1, x + 1))
-
-        self.assertFalse(flatMappedRDD.isCheckpointed())
-        self.assertTrue(flatMappedRDD.getCheckpointFile() is None)
-        self.assertFalse(self.sc.getCheckpointDir() is None)
-
-        flatMappedRDD.checkpoint()
-        result = flatMappedRDD.collect()
-        time.sleep(1)  # 1 second
-        self.assertTrue(flatMappedRDD.isCheckpointed())
-        self.assertEqual(flatMappedRDD.collect(), result)
-        self.assertEqual(
-            "file:" + self.checkpointDir.name,
-            os.path.dirname(os.path.dirname(flatMappedRDD.getCheckpointFile())),
-        )
-        self.assertEqual(
-            self.sc.getCheckpointDir(), os.path.dirname(flatMappedRDD.getCheckpointFile())
-        )
-
-    def test_checkpoint_and_restore(self):
-        parCollection = self.sc.parallelize([1, 2, 3, 4])
-        flatMappedRDD = parCollection.flatMap(lambda x: [x])
-
-        self.assertFalse(flatMappedRDD.isCheckpointed())
-        self.assertTrue(flatMappedRDD.getCheckpointFile() is None)
-
-        flatMappedRDD.checkpoint()
-        flatMappedRDD.count()  # forces a checkpoint to be computed
-        time.sleep(1)  # 1 second
-
-        self.assertTrue(flatMappedRDD.getCheckpointFile() is not None)
-        recovered = self.sc._checkpointFile(
-            flatMappedRDD.getCheckpointFile(), flatMappedRDD._jrdd_deserializer
-        )
-        self.assertEqual([1, 2, 3, 4], recovered.collect())
-
-
-class LocalCheckpointTests(ReusedPySparkTestCase):
-    def test_basic_localcheckpointing(self):
-        parCollection = self.sc.parallelize([1, 2, 3, 4])
-        flatMappedRDD = parCollection.flatMap(lambda x: range(1, x + 1))
-
-        self.assertFalse(flatMappedRDD.isCheckpointed())
-        self.assertFalse(flatMappedRDD.isLocallyCheckpointed())
-
-        flatMappedRDD.localCheckpoint()
-        result = flatMappedRDD.collect()
-        time.sleep(1)  # 1 second
-        self.assertTrue(flatMappedRDD.isCheckpointed())
-        self.assertTrue(flatMappedRDD.isLocallyCheckpointed())
-        self.assertEqual(flatMappedRDD.collect(), result)
-
-
-class AddFileTests(PySparkTestCase):
-    def test_add_py_file(self):
-        # To ensure that we're actually testing addPyFile's effects, check that
-        # this job fails due to `userlibrary` not being on the Python path:
-        # disable logging in log4j temporarily
-        def func(x):
-            from userlibrary import UserClass
-
-            return UserClass().hello()
-
-        with QuietTest(self.sc):
-            self.assertRaises(Exception, self.sc.parallelize(range(2)).map(func).first)
-
-        # Add the file, so the job should now succeed:
-        path = os.path.join(SPARK_HOME, "python/test_support/userlibrary.py")
-        self.sc.addPyFile(path)
-        res = self.sc.parallelize(range(2)).map(func).first()
-        self.assertEqual("Hello World!", res)
-
-    def test_add_file_locally(self):
-        path = os.path.join(SPARK_HOME, "python/test_support/hello/hello.txt")
-        self.sc.addFile(path)
-        download_path = SparkFiles.get("hello.txt")
-        self.assertNotEqual(path, download_path)
-        with open(download_path) as test_file:
-            self.assertEqual("Hello World!\n", test_file.readline())
-
-    def test_add_file_recursively_locally(self):
-        path = os.path.join(SPARK_HOME, "python/test_support/hello")
-        self.sc.addFile(path, True)
-        download_path = SparkFiles.get("hello")
-        self.assertNotEqual(path, download_path)
-        with open(download_path + "/hello.txt") as test_file:
-            self.assertEqual("Hello World!\n", test_file.readline())
-        with open(download_path + "/sub_hello/sub_hello.txt") as test_file:
-            self.assertEqual("Sub Hello World!\n", test_file.readline())
-
-    def test_add_py_file_locally(self):
-        # To ensure that we're actually testing addPyFile's effects, check that
-        # this fails due to `userlibrary` not being on the Python path:
-        def func():
-            from userlibrary import UserClass  # noqa: F401
-
-        self.assertRaises(ImportError, func)
-        path = os.path.join(SPARK_HOME, "python/test_support/userlibrary.py")
-        self.sc.addPyFile(path)
-        from userlibrary import UserClass
-
-        self.assertEqual("Hello World!", UserClass().hello())
-
-    def test_add_egg_file_locally(self):
-        # To ensure that we're actually testing addPyFile's effects, check that
-        # this fails due to `userlibrary` not being on the Python path:
-        def func():
-            from userlib import UserClass
-
-            UserClass()
-
-        self.assertRaises(ImportError, func)
-        path = os.path.join(SPARK_HOME, "python/test_support/userlib-0.1.zip")
-        self.sc.addPyFile(path)
-        from userlib import UserClass
-
-        self.assertEqual("Hello World from inside a package!", UserClass().hello())
-
-    def test_overwrite_system_module(self):
-        self.sc.addPyFile(os.path.join(SPARK_HOME, "python/test_support/SimpleHTTPServer.py"))
-
-        import SimpleHTTPServer
-
-        self.assertEqual("My Server", SimpleHTTPServer.__name__)
-
-        def func(x):
-            import SimpleHTTPServer
-
-            return SimpleHTTPServer.__name__
-
-        self.assertEqual(["My Server"], self.sc.parallelize(range(1)).map(func).collect())
-
-
-class ContextTests(unittest.TestCase):
-    def test_failed_sparkcontext_creation(self):
-        # Regression test for SPARK-1550
-        self.assertRaises(Exception, lambda: SparkContext("an-invalid-master-name"))
-
-    def test_get_or_create(self):
-        with SparkContext.getOrCreate() as sc:
-            self.assertTrue(SparkContext.getOrCreate() is sc)
-
-    def test_parallelize_eager_cleanup(self):
-        with SparkContext() as sc:
-            temp_files = os.listdir(sc._temp_dir)
-            sc.parallelize([0, 1, 2])
-            post_parallelize_temp_files = os.listdir(sc._temp_dir)
-            self.assertEqual(temp_files, post_parallelize_temp_files)
-
-    def test_set_conf(self):
-        # This is for an internal use case. When there is an existing SparkContext,
-        # SparkSession's builder needs to set configs into SparkContext's conf.
-        sc = SparkContext()
-        sc._conf.set("spark.test.SPARK16224", "SPARK16224")
-        self.assertEqual(sc._jsc.sc().conf().get("spark.test.SPARK16224"), "SPARK16224")
-        sc.stop()
-
-    def test_stop(self):
-        sc = SparkContext()
-        self.assertNotEqual(SparkContext._active_spark_context, None)
-        sc.stop()
-        self.assertEqual(SparkContext._active_spark_context, None)
-
-    def test_with(self):
-        with SparkContext():
-            self.assertNotEqual(SparkContext._active_spark_context, None)
-        self.assertEqual(SparkContext._active_spark_context, None)
-
-    def test_with_exception(self):
-        try:
-            with SparkContext():
-                self.assertNotEqual(SparkContext._active_spark_context, None)
-                raise RuntimeError()
-        except BaseException:
-            pass
-        self.assertEqual(SparkContext._active_spark_context, None)
-
-    def test_with_stop(self):
-        with SparkContext() as sc:
-            self.assertNotEqual(SparkContext._active_spark_context, None)
-            sc.stop()
-        self.assertEqual(SparkContext._active_spark_context, None)
-
-    def test_progress_api(self):
-        with SparkContext() as sc:
-            sc.setJobGroup("test_progress_api", "", True)
-            rdd = sc.parallelize(range(10)).map(lambda x: time.sleep(100))
-
-            def run():
-                # When thread is pinned, job group should be set for each thread for now.
-                # Local properties seem not being inherited like Scala side does.
-                if os.environ.get("PYSPARK_PIN_THREAD", "true").lower() == "true":
-                    sc.setJobGroup("test_progress_api", "", True)
-                try:
-                    rdd.count()
-                except Exception:
-                    pass
-
-            t = threading.Thread(target=run)
-            t.daemon = True
-            t.start()
-            # wait for scheduler to start
-            time.sleep(3)
-
-            tracker = sc.statusTracker()
-            jobIds = tracker.getJobIdsForGroup("test_progress_api")
-            self.assertEqual(1, len(jobIds))
-            job = tracker.getJobInfo(jobIds[0])
-            self.assertEqual(1, len(job.stageIds))
-            stage = tracker.getStageInfo(job.stageIds[0])
-            self.assertEqual(rdd.getNumPartitions(), stage.numTasks)
-
-            sc.cancelAllJobs()
-            t.join()
-            # wait for event listener to update the status
-            time.sleep(1)
-
-            job = tracker.getJobInfo(jobIds[0])
-            self.assertEqual("FAILED", job.status)
-            self.assertEqual([], tracker.getActiveJobsIds())
-            self.assertEqual([], tracker.getActiveStageIds())
-
-            sc.stop()
-
-    def test_startTime(self):
-        with SparkContext() as sc:
-            self.assertGreater(sc.startTime, 0)
-
-    def test_forbid_insecure_gateway(self):
-        # Fail immediately if you try to create a SparkContext
-        # with an insecure gateway
-        parameters = namedtuple("MockGatewayParameters", "auth_token")(None)
-        mock_insecure_gateway = namedtuple("MockJavaGateway", "gateway_parameters")(parameters)
-        with self.assertRaises(ValueError) as context:
-            SparkContext(gateway=mock_insecure_gateway)
-        self.assertIn("insecure Py4j gateway", str(context.exception))
-
-    def test_resources(self):
-        """Test the resources are empty by default."""
-        with SparkContext() as sc:
-            resources = sc.resources
-            self.assertEqual(len(resources), 0)
-
-    def test_disallow_to_create_spark_context_in_executors(self):
-        # SPARK-32160: SparkContext should not be created in executors.
-        with SparkContext("local-cluster[3, 1, 1024]") as sc:
-            with self.assertRaises(Exception) as context:
-                sc.range(2).foreach(lambda _: SparkContext())
-            self.assertIn(
-                "CONTEXT_ONLY_VALID_ON_DRIVER",
-                str(context.exception),
-            )
-
-    def test_allow_to_create_spark_context_in_executors(self):
-        # SPARK-32160: SparkContext can be created in executors if the config is set.
-
-        def create_spark_context():
-            conf = SparkConf().set("spark.executor.allowSparkContext", "true")
-            with SparkContext(conf=conf):
-                pass
-
-        with SparkContext("local-cluster[3, 1, 1024]") as sc:
-            sc.range(2).foreach(lambda _: create_spark_context())
-
-
-class ContextTestsWithResources(unittest.TestCase):
-    def setUp(self):
-        class_name = self.__class__.__name__
-        self.tempFile = tempfile.NamedTemporaryFile(delete=False)
-        self.tempFile.write(b'echo {\\"name\\": \\"gpu\\", \\"addresses\\": [\\"0\\"]}')
-        self.tempFile.close()
-        # create temporary directory for Worker resources coordination
-        self.tempdir = tempfile.NamedTemporaryFile(delete=False)
-        os.unlink(self.tempdir.name)
-        os.chmod(
-            self.tempFile.name,
-            stat.S_IRWXU | stat.S_IXGRP | stat.S_IRGRP | stat.S_IROTH | stat.S_IXOTH,
-        )
-        conf = SparkConf().set("spark.test.home", SPARK_HOME)
-        conf = conf.set("spark.driver.resource.gpu.amount", "1")
-        conf = conf.set("spark.driver.resource.gpu.discoveryScript", self.tempFile.name)
-        self.sc = SparkContext("local-cluster[2,1,1024]", class_name, conf=conf)
-
-    def test_resources(self):
-        """Test the resources are available."""
-        resources = self.sc.resources
-        self.assertEqual(len(resources), 1)
-        self.assertTrue("gpu" in resources)
-        self.assertEqual(resources["gpu"].name, "gpu")
-        self.assertEqual(resources["gpu"].addresses, ["0"])
-
-    def tearDown(self):
-        os.unlink(self.tempFile.name)
-        self.sc.stop()
-
-
-if __name__ == "__main__":
-    from pyspark.tests.test_context import *  # noqa: F401
-
-    try:
-        import xmlrunner
-
-        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
-    except ImportError:
-        testRunner = None
-    unittest.main(testRunner=testRunner, verbosity=2)

snowflake/snowpark_connect/includes/python/pyspark/tests/test_daemon.py

@@ -1,89 +0,0 @@
-#
-# Licensed to the Apache Software Foundation (ASF) under one or more
-# contributor license agreements.  See the NOTICE file distributed with
-# this work for additional information regarding copyright ownership.
-# The ASF licenses this file to You under the Apache License, Version 2.0
-# (the "License"); you may not use this file except in compliance with
-# the License.  You may obtain a copy of the License at
-#
-#    http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-#
-import os
-import sys
-import time
-import unittest
-
-from pyspark.serializers import read_int
-
-
-class DaemonTests(unittest.TestCase):
-    def connect(self, port):
-        from socket import socket, AF_INET, AF_INET6, SOCK_STREAM
-
-        family, host = AF_INET, "127.0.0.1"
-        if os.environ.get("SPARK_PREFER_IPV6", "false").lower() == "true":
-            family, host = AF_INET6, "::1"
-        sock = socket(family, SOCK_STREAM)
-        sock.connect((host, port))
-        # send a split index of -1 to shutdown the worker
-        sock.send(b"\xFF\xFF\xFF\xFF")
-        sock.close()
-        return True
-
-    def do_termination_test(self, terminator):
-        from subprocess import Popen, PIPE
-        from errno import ECONNREFUSED
-
-        # start daemon
-        daemon_path = os.path.join(os.path.dirname(__file__), "..", "daemon.py")
-        python_exec = sys.executable or os.environ.get("PYSPARK_PYTHON")
-        daemon = Popen([python_exec, daemon_path], stdin=PIPE, stdout=PIPE)
-
-        # read the port number
-        port = read_int(daemon.stdout)
-
-        # daemon should accept connections
-        self.assertTrue(self.connect(port))
-
-        # wait worker process spawned from daemon exit.
-        time.sleep(1)
-
-        # request shutdown
-        terminator(daemon)
-        time.sleep(1)
-
-        # daemon should no longer accept connections
-        try:
-            self.connect(port)
-        except EnvironmentError as exception:
-            self.assertEqual(exception.errno, ECONNREFUSED)
-        else:
-            self.fail("Expected EnvironmentError to be raised")
-
-    def test_termination_stdin(self):
-        """Ensure that daemon and workers terminate when stdin is closed."""
-        self.do_termination_test(lambda daemon: daemon.stdin.close())
-
-    def test_termination_sigterm(self):
-        """Ensure that daemon and workers terminate on SIGTERM."""
-        from signal import SIGTERM
-
-        self.do_termination_test(lambda daemon: os.kill(daemon.pid, SIGTERM))
-
-
-if __name__ == "__main__":
-    from pyspark.tests.test_daemon import *  # noqa: F401
-
-    try:
-        import xmlrunner
-
-        testRunner = xmlrunner.XMLTestRunner(output="target/test-reports", verbosity=2)
-    except ImportError:
-        testRunner = None
-    unittest.main(testRunner=testRunner, verbosity=2)