如何让Spark单元测试Mock代码在所有Worker节点进程生效

核心问题根因 #

• unittest.mock.patch是单进程内存级的对象替换，仅会修改当前Driver进程内的模块属性，不会同步到独立启动的Executor/Worker进程。哪怕是local[2]本地运行模式，Spark的Python Worker也是独立fork的子进程，启动时会重新加载所有依赖模块，完全不继承Driver进程的内存修改。
• 无Spark的单元测试能正常运行，本质是所有代码都在单Driver进程内执行，patch的内存修改可以正常生效；一旦提交任务到独立Worker进程，跨进程的内存同步断点就会直接触发原方法的文件读取逻辑，抛出找不到文件的错误。
• 当前业务代码的赋值时机存在硬编码触发问题：Env.user = user()是类定义加载阶段就会执行的逻辑，只要Worker进程导入spark_mock模块，就会立刻触发原user()方法读文件的操作，根本等不到Driver端的mock逻辑生效。
• 附带的测试代码笔误：test_env_without_spark中patch的目标路径是pkgname.spark_mock_dependency.hb_user，但实际模块中不存在该方法名，patch实际未生效。

推荐落地方案（全节点适配，低侵入） #

不要依赖跨进程的内存patch，改用通过Spark配置给所有节点统一下发mock标识的方式，从逻辑层面让Driver和所有Worker都走一致的mock分支，完全绕开进程间内存不同步的问题，具体步骤：

1. 修改原user()方法，增加统一的mock逻辑分支，优先读取节点上的环境变量：

import os

def user():
    # 优先读取全局mock标识，保证Driver和所有Worker逻辑完全一致
    mock_user = os.getenv("MOCK_USER")
    if mock_user:
        return mock_user
    # 保留原有业务逻辑
    with open('/etc/user') as f:
        return f.readline().strip()

2. 在测试类的SparkConf配置中，给所有Executor节点注入相同的环境变量，同时同步Driver进程的环境变量，避免两边逻辑分叉：

@classmethod
def setUpClass(cls) -> None:
    conf = SparkConf().setMaster("local[2]").setAppName("testing")
    # 给所有Executor节点下发统一的mock环境变量
    conf.set("spark.executorEnv.MOCK_USER", "anyone")
    # Driver进程同步设置相同的环境变量
    os.environ["MOCK_USER"] = "anyone"
    cls.sc = SparkContext(conf=conf)

3. 修正原测试中的遗留问题：
   • 把两个测试方法里的patch目标统一改成正确的pkgname.spark_mock_dependency.user（如果需要保留单进程下的patch逻辑的话）
   • 修正断言逻辑：原self.assertTrue([res == 'anyone' for res in results])永远为真（列表对象本身为布尔真值），改成精确匹配结果的断言：self.assertEqual(['anyone', 'anyone'], results)

临时调试方案（不推荐长期使用） #

如果短期不打算修改业务代码，需要保留原类定义加载即执行的逻辑，可以通过Worker启动钩子在每个节点初始化时自动打补丁：在SparkConf中配置Python Worker的初始化参数，让每个Worker进程启动时自动执行mock逻辑，替换内存中的user方法。但这种方式侵入性强，会污染所有Worker的运行环境，仅适合临时验证场景。

内容的提问来源于stack exchange，提问作者Andrew Stark