Python Unit test how to test a method which calls another method

Question

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?

Answer 1

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.