Python Unittest测试使用sys.argv定义全局变量脚本的正确方法

你遇到的报错本质是导入执行顺序问题：Python 导入my_script模块时，会立刻执行模块顶层的所有代码，也就是args = my_parser()这行逻辑会在测试用例的patch逻辑执行前就跑完，此时你还没替换sys.argv，自然会触发必填参数校验失败。

最简单的方式是延迟导入业务模块，等你完成sys.argv的 mock 之后再导入my_script相关内容，测试代码修改如下：

from unittest import TestCase, patch
import sys

class MyTest(TestCase):
  def test_do_something(self):
    test_sys_argv = ["dummy", "--my-arg", "dummy_arg"]
    with patch.object(sys, "argv", test_sys_argv):
      # 在patch上下文内导入my_script，此时模块顶层执行parser时读取的是mock后的参数
      from my_script import do_something
      out = do_something()
    self.assertEqual(out, 6)

如果有多个测试用例都需要用到my_script的内容，可以在测试文件最开头先完成全局 mock 再导入模块：

from unittest import patch
import sys

# 先mock sys.argv再导入业务模块
test_sys_argv = ["dummy", "--my-arg", "dummy_arg"]
patch.object(sys, "argv", test_sys_argv).start()

# mock完成后再导入
from unittest import TestCase
from my_script import do_something

class MyTest(TestCase):
  def test_do_something(self):
    out = do_something()
    self.assertEqual(out, 6)

如果可以对业务脚本做少量无侵入调整，可以给顶层参数解析逻辑加入口判断，仅在脚本直接运行时执行参数解析，导入时不触发，修改后的my_script.py如下：

import argparse, sys

def my_parser(args=sys.argv[1:]):
  parser = argparse.ArgumentParser(description="my parser")
  parser.add_argument("--my-arg", dest="my_arg", required=True)
  args, _ = parser.parse_known_args(args=args)
  return args

def do_something():
  res = 5+1
  return res

def main():
  use_args = args
  res = do_something()
  # whatevere here

# 仅直接运行脚本时执行参数解析和main函数
if __name__ == "__main__":
  args = my_parser()
  main()

调整后args全局变量在脚本直接运行时仍然可用，导入时不会触发参数校验，测试代码不需要做任何特殊处理就能直接运行。

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