python-3.x - 在databricks中，在不同笔记本中的函数上使用unittest.mock.patch

Question

在 databricks 上，我有一个代码笔记本和一个单元测试笔记本。

使用“%run”命令将代码“导入”到单元测试笔记本中。

如何从单元测试笔记本中创建代码笔记本中函数之一的模拟对象？我通常会使用补丁上下文管理器来实现此目的。

这是包含要修补的函数 (get_name) 的代码笔记本:

# Databricks notebook source
def get_name_func():
  return 'name1'

这是单元测试代码:

# Databricks notebook source:

from unittest.mock import patch
import inspect


# COMMAND ----------


# MAGIC %run ./get_name


# COMMAND ----------


def local_get_name():
  return 'name_local'


# COMMAND ----------


get_name_func()


# COMMAND ----------


print(inspect.getmodule(get_name_func))
print(inspect.getsourcefile(get_name_func))


# COMMAND ----------


inspect.unwrap(get_name_func)


# COMMAND ----------


with patch('get_name_func') as mock_func:
  print(mock_func)


# COMMAND ----------
with patch('local_get_name') as mock_func:
  print(mock_func)

对于本地函数和代码笔记本中的函数，两次修补尝试都会给出相同的错误:

TypeError: Need a valid target to patch. You supplied: 'get_name_func'

检查命令返回:

<module '__main__' from '/local_disk0/tmp/1625490167313-0/PythonShell.py'>
<command-6807918>

和

Out[38]: <function __main__.get_name_func()>

我尝试了模块路径的各种组合，但没有成功。

奇怪的是，__name__ 返回'__main__'。但是在补丁调用中使用路径 '__main__.get_name_func' 不起作用。

我的信念是，如果该对象存在于笔记本中(确实如此)，那么它一定是可修补的。

有什么建议吗？

Answer 1

我必须制作自己的修补功能:

class FunctionPatch():
  '''
  This class is a context manager that allows patching of functions "imported" from another notebook using %run.
  
  The patch function must be at global scope (i.e. top level)
  '''
  def __init__(self, real_func_name: str, patch_func: Callable):
    self._real_func_name = real_func_name
    self._patch_func = patch_func
    self._backup_real_func = None
    
  def __enter__(self):
    self._backup_real_func = globals()[self._real_func_name]
    globals()[self._real_func_name] = self._patch_func
    
  def __exit__(self, exc_type, exc_value, tb):
    if exc_type is not None:
      traceback.print_exception(exc_type, exc_value, tb)
      
    globals()[self._real_func_name] = self._backup_real_func

用法:

def test_function_patch_real_func():
  return 'real1'

def test_function_patch():
  
  assert test_function_patch_real_func() == 'real1'
  
  def mock_func():
    return 'mock1'
  
  with FunctionPatch('test_function_patch_real_func', mock_func):
    assert test_function_patch_real_func() == 'mock1' 
    
  assert test_function_patch_real_func() == 'real1'