PySpark单元测试：解决未关闭套接字等警告问题

咱们先逐个拆解这些警告的来龙去脉，再给出针对性的解决办法，让你的测试输出清爽起来！

1. ResourceWarning: unclosed <socket.socket [...]> / unclosed file <_io.BufferedWriter [...]> #

产生原因 #

这俩警告本质是一回事：PySpark的SparkSession或SparkContext在测试结束后没有被显式关闭，导致底层的网络socket连接、文件句柄等系统资源没被正常释放。Python的垃圾回收机制虽然最终会清理，但会检测到这些未主动关闭的资源，从而抛出警告。

尤其是在单元测试场景中，如果每个测试用例都创建新的Spark实例，又不做清理，就会频繁触发这类警告。

解决方法 #

最靠谱的方式是显式管理Spark实例的生命周期，推荐两种做法：

方法一：用测试类的setUpClass/tearDownClass统一管理

在测试类中，只创建一次SparkSession，所有测试用例共享，测试全部结束后再关闭：

from pyspark.sql import SparkSession
import unittest

class TestPySpark(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        # 创建本地模式的SparkSession，适合单元测试
        cls.spark = SparkSession.builder.master("local[1]").appName("UnitTest").getOrCreate()
    
    @classmethod
    def tearDownClass(cls):
        # 测试全部结束后关闭SparkSession
        cls.spark.stop()
    
    def test_insert_and_collect(self):
        df = self.spark.createDataFrame([(1,)], ["col"])
        result = df.collect()
        self.assertEqual(result[0][0], 1)

方法二：在函数内用with语句（上下文管理器）

如果必须在函数内创建Spark实例，用with语句自动关闭资源：

def insert_and_collect(value):
    with SparkSession.builder.appName("Test").getOrCreate() as spark:
        df = spark.createDataFrame([(value,)], ["col"])
        return df.collect()

with块结束后，SparkSession会自动调用stop()释放资源。

2. DeprecationWarning: invalid escape sequence #

产生原因 #

Python 3.6+对字符串中的未转义反斜杠做了更严格的检查：如果你的代码里出现了类似"C:\data\file"这种路径，或者正则表达式里的未转义反斜杠（比如"\d"），Python会认为这是无效的转义序列，抛出DeprecationWarning（未来版本会直接报错）。

解决方法 #

两种修复方式选其一：

• 用原始字符串：在字符串前加r，让Python忽略转义序列：

  path = r"C:\data\file"
  pattern = r"\d+"
  

• 用双反斜杠转义：把单个反斜杠改成两个：

  path = "C:\\data\\file"
  pattern = "\\d+"
  

临时方案：屏蔽警告（不推荐长期使用） #

如果暂时来不及从根源修复，或者确认这些警告不影响测试结果，可以在测试文件开头添加警告过滤：

import warnings
# 屏蔽ResourceWarning
warnings.filterwarnings("ignore", category=ResourceWarning)
# 屏蔽无效转义序列的DeprecationWarning
warnings.filterwarnings("ignore", category=DeprecationWarning, message="invalid escape sequence")

不过还是建议优先从根源解决问题，避免隐藏潜在的资源泄漏或代码兼容问题。

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