Pytest：创建 SparkSession

Posted 2023-03-23

【中文标题】Pytest：创建 SparkSession

【英文标题】：Pytest: create SparkSession

【发布时间】：2022-01-23 21:29:52

【问题描述】：

我需要使用 pytest 测试我的 Spark 项目，但我不明白如何创建 Spark 会话。我做了一些研究并想出了：

import pytest
import unittest
from pyspark.sql import SparkSession

@pytest.fixture(scope="session")
def spark_session():
    spark = SparkSession.builder.master("local[*]").appName("test").getOrCreate()
    return spark

class Testdb2connection(unittest.TestCase):
    @pytest.mark.usefixtures("spark_session")
    def test_connectdb2(self):
       with self.assertRaises(ValueError):
          return spark_session.format('jdbc')\
          ...

但是，在运行测试时，我得到：

'AttributeError: 'function' 对象没有属性'format'

我做错了什么？

【参考方案1】：

查看Mixing pytest fixtures into unittest.TestCase subclasses using marks，您可以定义spark_session 范围为class，并将火花会话添加到请求上下文的cls 属性中，以便能够将其用作使用该夹具的类中的属性。

尝试使用以下修改后的代码：

import pytest
import unittest
from pyspark.sql import SparkSession

@pytest.fixture(scope='class')
def spark_session(request):
    spark = SparkSession.builder.master("local[*]").appName("test").getOrCreate()

    request.addfinalizer(lambda: spark.stop()) # to teardown the session after test

    request.cls.spark = spark


@pytest.mark.usefixtures("spark_session")
class Testdb2connection(unittest.TestCase):

    def test_connectdb2(self):
        assert isinstance(self.spark, SparkSession)

运行测试：

pytest ./mytest.py

.                                        [100%]                                                          
1 passed in 4.12s

【讨论】：

非常感谢您的帮助！我确实可以使用您的代码成功运行相同的测试，但是在运行 spark_session.format('jdbc')\

时，'AttributeError: 'function' Object has no attribute 'format' 的一个问题仍然存在

@Moritz spark_session 夹具现在作为类属性传递，因此您需要使用 self.spark.read.format("jdbc")... 而不是 spark_session.format('jdbc')...