使用pytest/conftest.py时PySpark Fixture无法传入setup_method的问题

问题场景 #

我尝试创建测试数据集对Spark函数做单元测试（验证get_spark_table方法仅返回2行数据），但运行测试命令：

retcode = pytest.main([".", "-vv", "-k", "TestYourClass"])

时出现错误：

TypeError: setup_method() missing 1 required positional argument: 'spark'

现有代码 #

conftest.py

from pyspark.sql import SparkSession

@pytest.fixture(scope="session", name='spark')
def spark() -> SparkSession:
    spark = SparkSession.builder.appName("Testing").getOrCreate()
    yield spark

test_your_class.py

from pyspark.sql.types import StructType, StructField, StringType, IntegerType

class TestYourClass(): 
    """ 
    Test suite for YourClass functionality.
    """

    def setup_method(self, method, spark): 
        """ 
        Setup method that runs before each test method
        """

        py_data = [(1, "John"), (2, "Jane"), (3, "Bob")] 
        py_schema = StructType([
            StructField("id", IntegerType()), 
            StructField("name", StringType())]) 
        py_df = spark.createDataFrame(data = py_data, 
                                      schema = py_schema) 
        py_df.createOrReplaceTempView("people") 

    def get_spark_table(self, spark, tablename):
        query = f"SELECT * FROM {tablename} WHERE id > 1"     
        df = spark.sql(query) 
        return df 

    def test_get_spark_table(self, spark):
        """ 
        Test SQL query
        """ 
        df = self.get_spark_table(spark, tablename = "people") 
        assert df.count() == 2

问题原因 #

pytest的setup_method是固定生命周期方法，它的签名只能是def setup_method(self, method):，无法直接接收fixture参数。你在setup_method中添加spark参数，导致pytest调用时无法传递该参数，从而抛出类型错误。

解决方案 #

改用@pytest.fixture(autouse=True)装饰初始化方法，替代原生的setup_method，这样就能自动注入spark fixture，并且在每个测试方法前执行初始化逻辑。

修改后的test_your_class.py #

from pyspark.sql.types import StructType, StructField, StringType, IntegerType
import pytest

class TestYourClass(): 
    """ 
    Test suite for YourClass functionality.
    """

    @pytest.fixture(autouse=True)
    def setup(self, spark): 
        """ 
        自动执行的初始化逻辑，替代setup_method
        """
        py_data = [(1, "John"), (2, "Jane"), (3, "Bob")] 
        py_schema = StructType([
            StructField("id", IntegerType()), 
            StructField("name", StringType())]) 
        py_df = spark.createDataFrame(data = py_data, 
                                      schema = py_schema) 
        py_df.createOrReplaceTempView("people") 

    def get_spark_table(self, spark, tablename):
        query = f"SELECT * FROM {tablename} WHERE id > 1"     
        df = spark.sql(query) 
        return df 

    def test_get_spark_table(self, spark):
        """ 
        Test SQL query
        """ 
        df = self.get_spark_table(spark, tablename = "people") 
        assert df.count() == 2

说明 #

• @pytest.fixture(autouse=True)会让这个fixture自动在每个测试方法执行前运行，无需在测试方法中显式声明依赖
• 原有的conftest.py无需修改，spark fixture的定义是正确的
• 如果你不想使用autouse，也可以在测试方法中显式依赖这个setup fixture，比如：

  def test_get_spark_table(self, spark, setup):
      # 测试逻辑
  

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