How to make unit test mock code effect on every spark worker node(process)

user_15729729 0
Viewed 362

I have three files:

  1. spark_mock_dependency.py provides a user() method to read /etc/user,
  2. spark_mock.py for create an Env class, which uses user() method to get who is the user.
  3. spark_mock_test.py for the unit test of Env class.

I don't have /etc/user on my environment so I need to mock the method user() to fake it. However, the unit test test_env_without_spark works but test_env_with_spark does not. It looks like mocking only works on the driver node, and I cannot mock a class or methods over all the worker nodes(processes). See my code and the error below.

Is there anybody who knows how to mock methods over all the workers nodes(processes)?

spark_mock_dependency.py

def user():
    with open('/etc/user') as f:
        return f.readline().strip()

spark_mock.py

from pkgname.spark_mock_dependency import user


class Env:
    user = user()

spark_mock_test.py

import unittest
from unittest.mock import patch
from pyspark import SparkConf, SparkContext


class EnvTest(unittest.TestCase):
    sc = None

    @classmethod
    def setUpClass(cls) -> None:
        conf = SparkConf().setMaster("local[2]").setAppName("testing")
        cls.sc = SparkContext(conf=conf)

    @patch('pkgname.spark_mock_dependency.user')
    def test_env_with_spark(self, user_mocker):
        user_mocker.return_value = 'anyone'

        from pkgname.spark_mock import Env
        rdd = self.__class__.sc.parallelize([1, 2])
        results = rdd.map(lambda: f'{Env.user}').collect()
        self.assertTrue([res == 'anyone' for res in results])

    @patch('pkgname.spark_mock_dependency.hb_user')
    def test_env_without_spark(self, user_mocker):
        user_mocker.return_value = 'anyone'
        from pkgname.spark_mock import Env
        self.assertEqual('anyone', Env.user)

Error messages

py4j.protocol.Py4JJavaError: An error occurred while calling 
z:org.apache.spark.api.python.PythonRDD.collectAndServe.
: org.apache.spark.SparkException: Job aborted due to stage failure: Task 0 in stage 0.0 failed 1 times, most recent failure: Lost task 0.0 in stage 0.0 (TID 0) (executor driver): org.apache.spark.api.python.PythonException: Traceback (most recent call last):
"lib/python3.6/site-packages/pyspark/python/lib/pyspark.zip/pyspark/worker.py", line 586, in main
    func, profiler, deserializer, serializer = read_command(pickleSer, infile)
"lib/python3.6/site-packages/pyspark/python/lib/pyspark.zip/pyspark/worker.py", line 69, in read_command
    command = serializer._read_with_length(file)
"lib/python3.6/site-packages/pyspark/python/lib/pyspark.zip/pyspark/serializers.py", line 160, in _read_with_length
    return self.loads(obj)
"lib/python3.6/site-packages/pyspark/python/lib/pyspark.zip/pyspark/serializers.py", line 430, in loads
    return pickle.loads(obj, encoding=encoding)
"lib/python3.6/site-packages/pkgname/spark_mock.py", line 4, in <module>
    class Env:
"lib/python3.6/site-packages/pkgname/spark_mock.py", line 5, in Env
    user = user()
"lib/python3.6/site-packages/pkgname/spark_mock_dependency.py", line 2, in hb_user
    with open('/etc/user') as f:
FileNotFoundError: [Errno 2] No such file or directory: '/etc/user'
pyspark apache-spark mocking unit-testing python
2 Answers

You can import the user package in your unit test and mock that package itself

check below code:

def test_env_with_spark(self, user_mocker):

    from pkgname.spark_mock import Env, user
    user_mocker.patch("pkgname.spark_mock.user", return_value='anyone')
    rdd = self.__class__.sc.parallelize([1, 2])
    results = rdd.map(lambda: f'{Env.user}').collect()
user_107271440

Here's one possible approach when using pytest:

import other_module

@pytest.mark.parametrize("num_partitions", [2, 5])
def test_spark(monkeypatch, num_partitions):
    spark = pyspark.sql.SparkSession.builder.master(f"local[{num_partitions}]").getOrCreate()
    def f(index, iterator):
        for _ in iterator:
            pass
        monkeypatch.setattr(other_module.imported_module, "function", lambda: None)
        return []

    # Monkeypatch on each driver.
    spark.sparkContext.parallelize(list(range(num_partitions))).mapPartitionsWithIndex(f).collect()

    # Run actual test ...

    spark.stop()
user_11362080
Related