如何mock if __name__ == "__main__" 代码行以编写项目测试用例

if __name__ == "__main__" 的分支判断逻辑依赖模块运行时的__name__内置变量：模块被直接执行时__name__值为__main__，被作为依赖导入时值为模块自身的文件名，不需要使用常规mock工具做特殊模拟，可通过以下两种方案实现测试：

• 方案1：封装主逻辑为独立函数（推荐）

  对原有代码侵入性极低，可同时提升代码可复用性，是通用最佳实践
  原有代码示例：

  # demo.py
  if __name__ == "__main__":
      import sys
      input_val = sys.argv[1]
      print(f"执行逻辑，输入值为：{input_val}")
  

  简单重构后：

  # demo.py
  def main(args=None):
      import sys
      args = args or sys.argv[1:]
      input_val = args[0]
      print(f"执行逻辑，输入值为：{input_val}")
  
  if __name__ == "__main__":
      main()
  

  测试用例（pytest为例）：

  # test_demo.py
  from demo import main
  
  def test_main_logic(capsys):
      # 直接传入自定义参数调用主函数，不需要模拟__name__
      main(["测试输入"])
      output = capsys.readouterr().out.strip()
      assert output == "执行逻辑，输入值为：测试输入"
  

• 方案2：动态导入模块触发分支（无需修改原有代码）
  若无法修改原有模块的代码，可在测试时动态导入模块并指定__name__为__main__，触发对应分支执行：
  测试用例示例：

  # test_demo.py
  import importlib
  import sys
  import pytest
  
  def test_raw_main_branch(capsys, monkeypatch):
      # 模拟命令行参数
      monkeypatch.setattr(sys, "argv", ["demo.py", "测试输入"])
      # 动态导入模块时设置__name__为__main__
      spec = importlib.util.spec_from_file_location("__main__", "./demo.py")
      demo_module = importlib.util.module_from_spec(spec)
      spec.loader.exec_module(demo_module)
      # 校验执行结果
      output = capsys.readouterr().out.strip()
      assert output == "执行逻辑，输入值为：测试输入"
  

注意：不要直接修改已导入模块的__name__属性来触发分支，模块导入时顶层代码已经执行过一次，仅修改属性不会重新运行对应逻辑

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