如何判断一段Python代码是否在unittest下运行

Question

我有一个很大的项目，里面使用了Python的unittest模块来进行单元测试。

其中有一个小方法，它控制着系统行为的很多重要方面。我希望这个方法在进行单元测试时能返回一个固定的结果，这样测试的结果才会一致。不过，如果我每个单元测试都去模拟这个方法的行为，成本会很高。

有没有办法让我这个方法在单元测试时能够“感知”到自己正在被测试，从而在测试时改变它的行为呢？

Answer 1

你可以在程序运行时专门为测试修改一个函数。比如说：

module.py

def func():
    return random.randint()

test.py

import module

def replacement_func():
    return 4 # chosen by fair dice roll

module.func = replacement_func

# run unit tests here

现在，每当 module 里的代码调用 func() 时，它实际上会调用你定义的 replacement_func()。

Answer 2

我用 tox 设置了一个环境变量，方法是这样的：

[testenv]
setenv = TOX_TESTENV = true

然后在代码里，我检查这个变量是否已经被设置：

import os


if os.env.get('TOX_TESTENV'):
    # Code is running under test.
    # This is useful to configure log levels for example.

Answer 3

我对 unittest 模块了解不多，但如果你是直接运行文件来进行单元测试，可以用下面的 if 语句把测试代码包起来：

if __name__ == "__main__":

在这个 if 语句里的代码，只有在你直接运行这个模块的时候才会执行，而不是在其他地方被引入。根据文档，这就是你应该如何调用 unittest.main() 的方式。

https://docs.python.org/2/library/unittest.html

这假设你不是从命令行运行的。

编辑：你可以查看函数调用栈，试着找到 unittest.main() 函数。

import inspect

def in_unit_test():
  current_stack = inspect.stack()
  for stack_frame in current_stack:
    for program_line in stack_frame[4]:    # This element of the stack frame contains 
      if "unittest" in program_line:       # some contextual program lines
        return True
  return False

https://docs.python.org/2/library/inspect.html

这有点像是个小技巧，但 inspect 模块有很多有用的函数，可以帮助你查看代码的内部情况。

Answer 4

我的解决办法是在运行 unittest 之前，先设置一个环境变量 TEST_FLAG=true。比如：

TEST_FLAG=true python -m unittest discover -s tests -b

接下来，只需要检查这个变量是否被设置了。比如：

MONGODB_URI =
    os.environ.get('MONGODB_URI') if not os.environ.get('TEST_FLAG') 
        else os.environ.get('MONGODB_TEST_URI')

Answer 5

你可以检查一下，unittest模块是否已经加载。这个模块只有在运行测试的时候才会被加载。

>>> 'unittest' in sys.modules.keys()
False
>>> from unittest import TestCase
>>> 'unittest' in sys.modules.keys()
True