I'm using python's built-in unit test tool: unittest
class MyClass:
def a():
# some operations
self.spark = SparkSession.builder.config(xxxx).getOrCreate()
How can I test this method?
Since a()
calls SparkSession method directly.
How am I going to test it?
CodePudding user response:
Use unittest.mock.patch
to patch out SparkSession
. Here's a self-contained single-file example that follows your example code as closely as possible and produces a passing pytest
result:
SparkSession = None
class MyClass:
def a(self):
# some operations
self.spark = SparkSession.builder.config('xxxx').getOrCreate()
from unittest.mock import patch
@patch('test.SparkSession')
def test_myclass_a(mock_spark_session):
MyClass().a()
Note that if I didn't patch SparkSession
out, calling MyClass().a()
would raise an AttributeError
since SparkSession
is otherwise None
.
MyClass
is defined in a file I created called test.py
so I patched 'test.SparkSession'
; you should patch it in the module under test, not the test module.