Pytest：创建 SparkSession

Question

我需要使用 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'

我做错了什么？

Answer 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