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' Object has no attribute 'format'

我做错了什么？

Answer 1

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

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

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

Pytest：创建 SparkSession

Pytest: create SparkSession

python

pytest

apache-spark

pyspark