Как исправить условно импортированный модуль в pytest? ⇐ Python

[Anonymous]

У меня есть код Python, который работает локально и при запуске на Databricks. Для сохранения результатов используется другая функция в зависимости от того, где выполняется код.
В Databricks несколько вещей автоматически инициализируются при запуске записной книжки; среди них есть библиотека pyspark.
Поэтому, чтобы мой код работал локально, я импортирую pyspark следующим образом:

Код: Выделить всё

if "DATABRICKS_RUNTIME_VERSION" in os.environ:
from pyspark.sql import functions as F

def save_results_to_databricks(...):
# do stuff like F.col("relevant")

Теперь, при написании теста, я бы исправил F вот так:

Код: Выделить всё

class TestSaveResultsToDatabricks(unittest.TestCase):
@patch("path.to.module.F")
def test_save_results_to_databricks(self, MockFunctions):
# test stuff

Однако это приведет к ошибке:

AttributeError: не имеет атрибута «F»

Итак, как мне исправить функцию, которая недоступна локально?

Подробнее здесь: https://stackoverflow.com/questions/793 ... -in-pytest