单元测试无TensorFlow依赖的函数时能否mock TensorFlow导入？

完全可以通过几种成熟方案跳过不必要的TF加载，大幅提升非TF依赖函数的测试速度，按需选择即可：

• 提前注入Mock模块（无需修改业务代码，通用首选）
  Python导入模块时会优先读取sys.modules里已缓存的模块对象，只要在导入被测业务文件之前，把TensorFlow相关的模块替换成Mock对象，就不会触发真实的TF加载逻辑。
  测试文件最开头加入以下代码即可，注意这段代码必须放在所有业务模块导入语句之前：

  import sys
  from unittest.mock import MagicMock
  
  # 注入TF主模块Mock
  sys.modules["tensorflow"] = MagicMock()
  # 如果业务代码导入了TF子模块，按需补充对应Mock即可
  sys.modules["tensorflow.keras"] = MagicMock()
  sys.modules["tensorflow.data"] = MagicMock()
  
  # 之后再导入被测模块，不会触发真实TF加载
  from your_business_module import non_tf_dependent_func
  

  如果测试逻辑里用到了TF的特定属性、常量，直接给Mock对象赋值对应测试值就行，不需要加载真实TF。

• 环境变量开关控制导入（性能最优，需少量修改业务代码）
  如果你可以修改业务代码，直接给TF导入逻辑加环境变量判断，测试时直接跳过整个导入分支，连Mock的开销都没有：
  业务文件中修改导入逻辑：

  import os
  if not os.getenv("UNIT_TEST_SKIP_TF"):
      import tensorflow as tf
  else:
      tf = None  # 非TF依赖的函数不会调用该对象，不影响逻辑
  

  跑单元测试时带上环境变量启动即可，比如pytest启动命令：

  UNIT_TEST_SKIP_TF=1 pytest tests/
  

  注意这个方案要确保你测试的函数完全不走到TF调用分支，否则会触发NoneType报错。

• Pytest全局自动注入（适配Pytest测试框架，无重复代码）
  如果你用Pytest做测试框架，直接在测试目录的conftest.py里加启动钩子，测试会话启动时自动完成TF Mock，不用每个测试文件重复写注入逻辑：

  import sys
  from unittest.mock import MagicMock
  
  def pytest_configure(config):
      # Pytest加载测试用例前完成Mock注入
      sys.modules["tensorflow"] = MagicMock()
      # 按需补充你用到的TF子模块Mock
      for sub_module in ["keras", "train", "saved_model"]:
          sys.modules[f"tensorflow.{sub_module}"] = MagicMock()
  

注意事项：如果你的测试集里同时包含TF依赖的用例和非TF用例，不要用全局Mock或者全局环境变量的方案，把两类用例拆到不同测试目录/文件分开执行即可，避免Mock影响真实TF逻辑的测试。

内容的提问来源于stack exchange，提问作者Rafael Dias