Python 单元测试

Table of Contents

1. 测试驱动开发

先写单元测试，然后再实现功能，这样做的好处：

  • 在长期实施测试驱动开发，可以明显感受到养成了良好的编码风格，比如函数与类的单一职责更是自然行成；
  • 修补 bug 和重构代码更加稳固。尤其是在代码变多的情况下，没有单元测试就会经常出现改好这里，又引出新的问题，测试用例可以更好的避免这种情况；
  • 在我的实践中，安全相关的关键代码也写入测试用例，保证源头的基本安全性。review 安全相关的测试用例时也能发现哪些安全问题没有覆盖到。举个例子，假如实现一个文件读取函数，我的单元测试会包含正常读取文件、越目录读取的测试用例。

整个实践的过程很简单，反复循环执行三个步骤：

  1. 测试先行，站在程序员使用的角度，写好功能测试；
  2. 运行测试用例，失败以后再写功能代码，直到测试用例通过；
  3. 检查代码是否值得继续重构。

编写测试用例、实现功能、重构代码，满足了这三大块，才叫TDD。

重点：

  1. 功能逐个实现，切勿一次写多个测试用例，再写代码。
  2. 重构时，先写测试用例，不要修改测试用例的同时也修改功能代码。
  3. 单元测试是用来测试流程和代码逻辑的，无关的东西不要测试（比如常量）。

2. 工具

2.1. pytest

官方网址：http://pytest.org/

在项目中，应当有良好的目录结构，通常把所有的单元测试文件放到单独的 test 目录中： 

|-- app/
  -- app.py
|-- test/
  -- test_app.py

2.2. 测试覆盖率

Coverage.py 用于统计单元测试的测试覆盖率的工具。

运行单个测试用例脚本： 

$ coverage run test/test_app_browser.py

然后生成覆盖率统计报告： 

$ coverage report
Name                                                                                                               Stmts   Miss  Cover
--------------------------------------------------------------------------------------------------------------------------------------
app/__init__.py                                                                                                       47      0   100%
app/api.py                                                                                                           195    150    23%
app/browser/__init__.py                                                                                                0      0   100%

输出的各列含义：

Stmts：有效代码总行数（没有算注释、空行）

Miss：未执行的代码总行数（没有算注释、空行）

Cover：覆盖率

如果用的 pytest，建议使用 pytest-cov 这个工具，它封装了 Coverage。使用方式也简单： 

$ py.test -v --cov=app 单元测试目录

注意 –cov 参数，它用于过滤所属项目的测试结果，如果没有过滤的话，每个执行过的模块都被打印出来。

3. doctest

doctest 会去执行 docstring 中的 Python 交互式会话中的内容，例如： 

# filename: common.py

import hashlib


def md5(s):
    """
    >>> md5("123")
    '202cb962ac59075b964b07152d234b70'
    >>> md5(123)
    Traceback (most recent call last):
    ...
    AssertionError
    """
    assert(s.__class__ == str)
    return hashlib.md5(s.encode("utf-8")).hexdigest()

执行测试命令： 

$ python3 -mdoctest common.py

如果未通过测试，会给出详细的提示。

doctest 更重要的是起调用示例的作用——给别人看的，或是生成自动文档。

4. 基本的断言——assert

用 assert 可以写出最简单的单元测试，比如 common.py 定义了一些项目中公用的函数： 

import hashlib


def md5(s):
    assert(s.__class__ == str)
    return hashlib.md5(s.encode("utf-8")).hexdigest()


def flat_list(a_list):
    result = []

    for l in a_list:
        result.extend(l)

    return result

现在在 test_common.py 中写一些测试用例： 

def test_md5():
    from common import md5
    assert(md5("123") == "202cb962ac59075b964b07152d234b70")


def test_flat_list():
    from common import flat_list
    assert(flat_list([]) == [])
    assert(flat_list([[1, 2], [3, 4]]) == [1, 2, 3, 4])

然后用 pytest 运行测试用例： 

$ pytest test_common.py

不建议在测试用例中直接使用 assert 关键字：

  • assert 不会显示出更详细的信息，仅知道触发了 AssertionError 异常。
  • 在运行某几个测试用例以前需要做一些环境初始化的工作。
  • 需要测试触发的异常。

5. unittest库

Python 标准库自带的单元测试库。

以前面的 common.py 为例，现在改用 unittest 编写单元测试代码： 

import unittest


class TestCommon(unittest.TestCase):
    def test_md5(self):
        from common import md5
        self.assertEqual(md5("123"), "202cb962ac59075b964b07152d234b70")
        self.assertRaises(AssertionError, md5, 123)  # 捕获抛出异常，如果为 md5 函数提供数字型就抛出异常

    def test_flat_list(self):
        from common import flat_list
        self.assertRaises(TypeError, flat_list, [1, 2, 3, 4])
        self.assertEqual(flat_list([]), [])
        self.assertEqual(flat_list([[1, 2], [3, 4]]), [1, 2, 3, 4])


if __name__ == '__main__':
    unittest.main()

说明：

  • 测试用的类以 Test 开头命名
  • 测试类继承 unittest.TestCase 类
  • 每个测试用例为一个方法，以 test_ 开头
  • 执行 unittest.main() 可以启动自动测试

更详细的用法请见官方文档：https://docs.python.org/3/library/unittest.html

5.1. setUp 和 tearDown

有一些测试场景比较复杂，在运行一组测试之前，需要依赖上下文的初始化，比如：

  • 需要先初始化数据库
  • 测试 API 接口，需要先获得 token

等等。

这时就可以借助自己实现 setUp 和 tearDown 方法。setUp 和 tearDown 跟类的构造函数和析构函数是一样的，setUp 用于测试前做环境初始化，如要测试数据库，则在 setUp 中需要先行连接数据库；tearDown 则是销毁环境作用，如关闭数据库。

如下演示的代码，被测试的接口需要先完成“登录”授权，我把登录授权和销毁状态功能实现在 setUp 和 tearDown 中： 

# 注：self.client、self.context 就定义在 TestBase 类中的
class TestManagerBlog(TestBase, unittest.TestCase):
    def setUp(self):
        # 先设置登录状态
        with self.client.session_transaction() as s:
            s["username"] = "test"
            s["islogin"] = True

    def tearDown(self):
        # 测试完成后销毁登录状态
        with self.client.session_transaction() as s:
            s.pop("username")
            s.pop("islogin")

    def test_post_new_blog(self):
        """发表文章测试"""
        with self.context:
            resp = self.client.post(url_for("blog.new_blog"),
                                    data={
                                        "title": "just test",
                                        "content": "test",
                                    })

            # 成功增加博文后，服务端返回 302，跳转到首页
            # 模拟 POST 请求后，判断状态码
            self.assertTrue(resp.status_code == 302)
            # 判断是否重定向首页
            self.assertTrue(resp.location.endswith("/"))

            # 查询数据库，看是否新增成功
            # 使用的 SQLAlchemy 库
            blog = ManagerBlogModel.query.filter(
                ManagerBlogModel.title == "just test"
            ).first()

            self.assertTrue(blog is not None)

6. Stub（桩代码） 测试

比如要测试的有些类，它同时依赖其他类，我们可以自己去实现一个和它依赖的类长得差不多的类。

示例代码： 

class Reader(object):
    def __init__(self, path):
        self.path = path

    def read(self):
        content = None

        try:
            with open(self.path) as f:
                content = f.read().strip()
        except FileNotFoundError:
            pass

        return content


class Flag(object):
    """验证 flag 文件"""
    def __init__(self):
        self.reader = Reader("/tmp/flag")

    def is_flag(self):
        flag = self.reader.read()
        if flag is None:
            raise ValueError("flag 错误")

        if flag == "49f68a5c8493ec2c0bf489821c21fc3b":
            return True

        return False

现在要测试 Flag 这个类，但要测试的对象类依赖 Reader 类，因此我们在测试代码中直接制造假的 Reader 类： 

import unittest
from test import Flag
from test import Reader


class StubReader1(Reader):
    def read(self):
        return "49f68a5c8493ec2c0bf489821c21fc3b"


class StubReader2(Reader):
    def read(self):
        return "xxx"


class StubReader3(Reader):
    def read(self):
        return None


class TestFlag(unittest.TestCase):
    def setUp(self):
        self.flag = Flag()

    def test_is_flag(self):
        # 测试前先修改依赖
        self.flag.reader = StubReader1("/tmp/flag")
        self.assertTrue(self.flag.is_flag())

        self.flag.reader = StubReader2("/tmp/flag")
        self.assertFalse(self.flag.is_flag())

        self.flag.reader = StubReader3("/tmp/flag")
        self.assertRaises(ValueError, self.flag.is_flag)

7. Mock 测试

有些测试是不依赖环境的（比如某个算法类的功能），但在实际项目中，有很功能依赖网络环境、读取外部文件等等，这些外部环境出现问题就会导致单元测试过不了，而用 Mock 测试就可以直接“模拟”出这些“环境”，保证每次测试都是如期的。

Mock 和 Stub 的区别： 

Stub 虽然是虚拟出的对象，但这个“虚拟”却要自己去实现一大堆代码了，比如要测试某个类的某个方法，需要先写一个继承类，然后覆盖父类的方法。Stub 有具体实现的方法，所以写测试时，还要关注它的实现逻辑。

对于 Mock，我们在不去重新实现依赖方法的情况下，就可以去做更深层次的依赖模拟，同时还能知道依赖的对象最终是否被调用到。

7.1. mock 库

Python 3.3+ 已经自带 mock 库（unittest.mock），如果使用的低版本 Python，可用 pip 安装： 

$ sudo pip install mock

mock 详细文档请见：https://docs.python.org/3/library/unittest.mock.html

示例1： 

from unittest import TestCase
from unittest.mock import patch

class TestCountIP(TestCase):
    def setUp(self):
        self.c = CountIP()

    @patch("config.es.search", lambda **args: {
        "took" : 26,
        "timed_out" : False,
        "_shards" : {
            "total" : 130,
            "successful" : 130,
            "skipped" : 0,
            "failed" : 0
        },
        "hits" : {
            "total" : 27073841,
            "max_score" : 0.0,
            "hits" : [ ]
        },
        "aggregations" : {
            "attack_ip_total" : {
                "value" : 11210866
            }
        }
    })
    def test_count(self):
        self.assertEqual(self.c.count(), 11210866)

以上代码中，CountIP 类中有个 count 方法负责调用 Elasticsearch 接口，然后取值。这个测试用例用 patch 装饰器指定好了 config.es.search 的调用结果，最后再做测试。

示例2，如何 mock datetime.datetime 模块

比如有一个获取昨日日期的函数，需要做单元测试： 

import datetime


def yesterday():
    now = datetime.datetime.now()
    yesterday = now - datetime.timedelta(days=1)
    return yesterday.strftime("%Y%m%d")

如果单元测试这样写： 

import datetime
from utils import yesterday
from unittest import TestCase, main
from unittest.mock import patch


class TestUtils(TestCase):
    @patch("datetime.datetime.now",
           lambda: datetime.datetime(2019, 11, 11, 11, 19, 26, 469747))
    def test_yesterday(self):
        self.assertEqual(yesterday(), "20191110")

执行测试时会报错： 

Traceback (most recent call last):
  File "/usr/lib64/python3.7/unittest/mock.py", line 1268, in patched
    patching.__exit__(*exc_info)
  File "/usr/lib64/python3.7/unittest/mock.py", line 1430, in __exit__
    setattr(self.target, self.attribute, self.temp_original)
TypeError: can't set attributes of built-in/extension type 'datetime.datetime'

对于这种情况，可以用匿名类来解决： 

class TestUtils(TestCase):
    @patch("datetime.datetime",
           type("XX", (datetime.datetime,), {
               "now": lambda: datetime.datetime(2019, 11, 11, 11, 19, 26, 469747)}))
    def test_yesterday(self):
        self.assertEqual(yesterday(), "20191110")

7.2. pytest mock

用法可参考文档：http://pytest.org/latest/monkeypatch.html

假如 common.py 中定义了个获取当前日期的函数： 

import time

def current_date():
    return time.strftime('%Y-%m-%d', time.localtime(time.time()))

想写个测试用例来判断它是否以“yyyy-MM-dd”格式返回当前日期，按 pytest mock 的文档写的 mock 测试： 

def test_current_date(monkeypatch):
    import time
    from common import current_date
    monkeypatch.setattr(time, "time", lambda: 1490749983.466406)
    assert(current_date() == "2017-03-29")

运行 pytest： 

$ pytest test_common.py

7.3. 其他 mock 库

对于一些特殊的三方库，差不多都可以在网上找到相应的专用 mock 库，在此展示两个。

7.3.1. mockredis

官网：https://github.com/locationlabs/mockredis

在某个 Flask 开发的 Web 应用中，用到了 Redis 做缓存，并实现了两个缓存读写函数，对读写函数的测试用例如下： 

"""缓存模块单元测试"""

from unittest.mock import patch
from unittest import TestCase
from mockredis import MockRedis

from flask import Flask
from flask_redis import Redis

from cache import write_cache
from cache import read_cache

@patch('cache.redis_conn', MockRedis(strict=True))
class TestCache(TestCase):
    def setUp(self):
        redis_conn = Redis()
        self.app = Flask(__name__)
        redis_conn.init_app(self.app)

    def test_write_cache(self):
        write_cache("test_redis", "test", second=30)
        self.assertTrue(read_cache("test_redis") == b"test")

        # 如果 second 参数为0，不做缓存
        write_cache("test_redis_0", "test", second=0)
        self.assertIsNone(read_cache("test_redis_0"))

    def test_read_cache(self):
        self.assertIsNone(read_cache("test"))

7.3.2. requests-mock

如果被测试的对象用到 Requests 库，就可以到 requests-mock，官网：https://requests-mock.readthedocs.io/

例，测试 API 登录获得 token 功能： 

import requests


def get_token(username, password):
    resp = requests.post("http://api.server.com/login",
                         data={"username": username, "password": password})
    return resp.text

测试用例如下： 

import requests_mock


def test_login():
    with requests_mock.Mocker() as m:
        m.post('http://api.server.com/login', text='{"token": "123"}')
        token = get_token(username="admin", password="admin")
        assert(token == '{"token": "123"}')

8. 小脚本中的单元测试

工作中经常会遇到只用写一两个文件的程序，不想搞复杂，通常我会在程序中直接定义一个测试函数去测试各项功能。然后为程序顺带提供一个 --test 参数，如： 

def self_test():
    ...


if __name__ == '__main__':
    ...
    parser.add_option("--test", help=u"自测", action="store_true")
    ...

    (options, args) = parser.parse_args()

    if options.test:
        self_test()
        raise SystemExit()