PySpark UDF中requests调用Mock失效问题排查与解决

问题根源 #

Spark UDF是分布式执行的：你在Driver端用unittest.mock Mock的src.my_funcs.requests.get仅作用于Driver进程，而UDF实际会被分发到Executor进程中执行。这些Executor进程会重新加载你的模块，完全不受Driver端Mock的影响，所以依然会调用真实的requests.get，导致URL格式错误的异常。

可行解决方案 #

方法一：拆分测试逻辑，先验函数再验UDF执行 #

这是最直接的测试方案，分两步验证：

1. 单独测试核心函数：去掉UDF装饰器，直接测试fetch_details函数，此时Mock可以正常生效，先确保逻辑正确。
2. 在本地单进程模式下测试UDF：配置SparkSession为local[1]（单进程模式），调用UDF后用collect()将结果拉回Driver，此时UDF会在Driver端执行，Mock能覆盖到。

测试代码示例：

import unittest
from unittest.mock import patch
from pyspark.sql import SparkSession
from src.my_funcs.udf_module import fetch_details_udf, fetch_details

class TestFetchUDF(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # 初始化本地单进程SparkSession
        cls.spark = SparkSession.builder \
            .master("local[1]") \
            .appName("TestUDF") \
            .getOrCreate()

    @classmethod
    def tearDownClass(cls):
        cls.spark.stop()

    # 测试核心函数逻辑
    def test_fetch_details_logic(self):
        with patch("src.my_funcs.udf_module.requests.get") as mock_get:
            mock_get.return_value.json.return_value = {"id": 1, "name": "test"}
            result = fetch_details("http://example.com")
            self.assertEqual(result, {"id": 1, "name": "test"})
            mock_get.assert_called_once_with("http://example.com")

    # 测试UDF在DataFrame中的行为
    def test_udf_in_dataframe(self):
        test_data = [("http://example.com",)]
        df = self.spark.createDataFrame(test_data, ["url"])
        
        with patch("src.my_funcs.udf_module.requests.get") as mock_get:
            mock_get.return_value.json.return_value = {"id": 1, "name": "test"}
            result_df = df.withColumn("details", fetch_details_udf("url"))
            # collect()触发UDF执行，结果拉回Driver
            results = result_df.collect()
            self.assertEqual(results[0]["details"], {"id": 1, "name": "test"})
            mock_get.assert_called_once_with("http://example.com")

方法二：修改UDF代码，动态导入requests #

将UDF内部的requests导入从模块顶部移到函数内部，这样Executor进程在执行UDF时才会加载requests模块。结合本地模式测试，Mock可以覆盖到Executor的执行环境：

修改后的UDF代码：

# src/my_funcs/udf_module.py
from pyspark.sql.functions import udf

def fetch_details(url):
    # 函数内部动态导入requests
    import requests
    resp = requests.get(url)
    return resp.json()

fetch_details_udf = udf(fetch_details)

测试时依然使用local[1]模式，Mock写法和方法一一致即可。

方法三：强制UDF在Driver端执行 #

通过Spark配置强制小数据量的UDF在Driver端运行，避免分发到Executor：

# 在测试的SparkSession中添加配置
spark.conf.set("spark.sql.execution.arrow.pyspark.enabled", "true")
spark.conf.set("spark.sql.execution.local", "true")

这种方式下，UDF会直接在Driver端执行，Mock可以正常生效。

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