Python unittest mock 小记

自从入坑 TDD 至今,并没有太多的 mock 需求,然而在随着深入到微服务后端组件的开发,依赖的组件越来越多也越来越诡异,便开始真正的正视起 Mock 来。

mock 是 Python 中一个神奇的模块,能够在测试中动态的替换部分逻辑,可以用来解决在跑测试的时候对其他组件依赖的问题。

在 Python 3.3 之后,mock 被加入到 unittest 模块里,不需要单独而在,而在之前,需要单独安装:

pip install mock

对于 mock 模块的学习也非常简单,只需要了解 Mock 类的使用和 Patch 的使用就非常受用了,前者是 mock 的精妙所在,可以创造一个可以替换任意函数、类、对象、方法的 object,并让这个 object 伪装成理想的样子。而 Patch 即是将这个 Mock object 替换的它应该存在的地方。

Mock

要理解 Mock,必须要先理解 Mock 类:

class unittest.mock.Mock(spec=None, side_effect=None, return_value=DEFAULT, wraps=None, name=None, spec_set=None, unsafe=False, **kwargs)

文档描述很详尽,但是常用的参数 side_effect, 和 return_value 值得注意。

side_effect 接受一个函数或异常,当 mock 实例被调用是,便会返回这个函数的返回内容,或者抛出对应的异常:

In [3]: m = Mock(side_effect=lambda: "hello")

In [4]: m()
Out[4]: 'hello'

In [5]: m = Mock(side_effect=Exception("Boooooom"))

In [6]: m()
---------------------------------------------------------------------------
Exception                                 Traceback (most recent call last)
<ipython-input-6-7e5925b669a0> in <module>()
----> 1 m()

....

Exception: Boooooom

当对于返回值比较固定的时候,return_value 可能会更加的实用:

In [7]: m = Mock(return_value="hello")

In [8]: m()
Out[8]: 'hello'

除了在 init 的时候对这些值进行注入,还可以通过赋值的方式:

In [10]: m.func.return_value = "world"

In [11]: m.func()
Out[11]: 'world'

目前只是对 Mock 的简单用法,因为 Mock 的实现很神奇,可以直接调用任何参数、方法,哪怕这些参数、方法并不存在,Mock 实例会直接返回一个新的 Mock 实例,像一棵树一样,非常有趣:

In [13]: m.wat
Out[13]: <Mock name='mock.wat' id='4406809544'>

In [14]: m.wat()
Out[14]: <Mock name='mock.wat()' id='4406812456'>

正是有了这样的特性,我们可以直接进行更深入的修改:

In [17]: m.wat.func.return_value = "watman"

In [18]: m.wat.func()
Out[18]: 'watman'

自此,如何 Mock 类已经非常清晰了,一个 Mock 实例可以作为 callable 的对象存在,而它本身,却又可以包含"无限"层级的新的的 Mock 实例,因此他可以非常简单的伪造出复杂的结构,比如 Response 等。

然而,在 mock 模块中,提供的并不只有一个 Mock 类,以上只是最基本的,还有更自动和常用的 unittest.mock.MagicMock,以及可以替换属性的 unittest.mock.PropertyMock。(除此之外还有两种 NonCallableMock)

MagicMock

MagicMock 是 Mock 的子类,其区别正如其名,即是实现了很多 magic method,可以举一个简单的例子,对于一个实例,如果没有实现 __iter__ 方法时,被转换成 list 会抛出异常:

In [19]: m = Mock()

In [20]: list(m)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-20-150120f9fb24> in <module>()
----> 1 list(m)

TypeError: 'Mock' object is not iterable

如果想正常的转化成列表,那么需要实现这个方法:

In [24]: m.__iter__ = Mock(return_value=iter([]))

In [25]: list(m)
Out[25]: []

MagicMock 类则会自作聪明的 mock 掉所有的 magic method,所以不需要自己去实现,而 __iter__ 默认是空数组:

In [28]: m = MagicMock()

In [29]: list(m)
Out[29]: []

我之所以认为 MagicMock 是自作聪明是因为如果我们需要这些魔法方法,默认的肯定不足以我们使用,我们无法避免自己重写。但因为已经有默认的存在,以至于如果我们手动实现的 Mock 实例存在问题时,往往会被它的默认方法掩盖掉一部分真相。

PropertyMock

PropertyMock 是我又爱又恨的 Mock 类,爱其有用,恨其难写。

在 Mock 类中默认 mock 的是可调用对象,因此,如果是熟悉或者 property 属性的 mock 并不理想:

In [30]: m = MagicMock()

In [31]: m.property_attr
Out[31]: <MagicMock name='mock.property_attr' id='4407739728'>

因此才需要 PropertyMock 的存在来拯救世界:

In [34]: p = PropertyMock(return_value=3)

In [35]: type(m).property_attr = p

In [36]: m.property_attr
Out[36]: 3

这种写法我真是日了狗了,不做详解,荆轲刺秦王。

补充

最后,作为 Mock 类的补充。Mock 除了提供了伪造数据的功能,还保留了大量的自带断言,可以来 ”事后“ 审计被 mock 的方法调用的情况,在单元测试中,简直神器。

Patch

有强大的 Mock 并不够用,只有真正的把伪造的可调用对象替换进去才算完美,而 Patch 便是负责这个工作。

在讲 mock 模块中实现的各种神奇 Patch 之前,依然先提下最基础的 unittest.mock.patch 类:

unittest.mock.patch(target, new=DEFAULT, spec=None, create=False, spec_set=None, autospec=None, new_callable=None, **kwargs)

一目了然,其中 target 是必选项,即使要替换什么,在基础的 Patch 类中,要求为字符串,比如:package.module.ClassName,而 new 即使要替换成的 Mock 类实例,默认是一个 MagicMock 实例,正如前面所说,这个实例实现了大量的默认魔法方法,只能能够让被调用时不抛异常而已,如果需要有数据替换等更精细的操作,需要自己定义 Mock 实例。

Patch 的参数基本上都是作为实例化 Mock 实例时使用的,如果自己手动构造 Mock 实例的话,只需要专注于 Patch 本身的功能就好。可以看一个简单的例子:

In [37]: class Class:
    ...:     def method(self):
    ...:         pass
    ...:

In [38]: with patch('__main__.Class') as MockClass:
    ...:     instance = MockClass.return_value
    ...:     instance.method.return_value = 'foo'
    ...:     assert Class() is instance
    ...:     assert Class().method() == 'foo'

值得一提的是,Patch 可以作为装饰器或者上下文管理器,前者对于默认的 Mock 实例即可满足需求的场景非常实用,而且 Patch 的额外参数也显得非常好用,后者对于构造负责的 Mock 实例非常好用。

Patch 拓展

Patch 类还有三个非常有用的拓展,分别是替换字典用的 patch.dict,多次替换用的 patch.multiple,和替换实例用的 patch.object

前两个非常简单,patch.dict 可以替换字典,或者字典类似实例,参考文档上的例子,一目了然:

>>> foo = {}
>>> with patch.dict(foo, {'newkey': 'newvalue'}):
...     assert foo == {'newkey': 'newvalue'}
...
>>> assert foo == {}

对于 patch.multiple,个人感觉有点鸡肋,目前也没真正的用过,作用正如文档所说:

Perform multiple patches in a single call.

文档的例子:

>>> thing = object()
>>> other = object()
>>> @patch.multiple('__main__', thing=DEFAULT, other=DEFAULT)
... def test_function(thing, other):
...     assert isinstance(thing, MagicMock)
...     assert isinstance(other, MagicMock)
...
>>> test_function()

patch.object 是需要单独拿出来强调的,毕竟和前两者相比,替换一个实例的场景是非常场景的,而且会有非常诡异的替换方式。

其实可以看到,Patch 的参数大同小异,我们只需要关注差异就可以了:

patch.object(target, attribute, new=DEFAULT, spec=None, create=False, spec_set=None, autospec=None, new_callable=None, **kwargs)

首先依然有 target,并且多出了 attribute 参数。对于 Patch 来说,要求 target 是字符串,而这里则要求是要替换的类,反而 attribute 被要求为字符串。除了这点,方法没有其他要注意的了,直接举一个文档的例子:

>>> @patch.object(SomeClass, 'class_method')
... def test(mock_method):
...     SomeClass.class_method(3)
...     mock_method.assert_called_with(3)
...
>>> test()

对于 Python unittest mock 的基础用法算是写完了,可以看到 Mock 和 Patch 都是非常灵性的工具,虽然把实现的奇技淫巧都跟封装了起来,但是其灵活性依然让人惊讶。至于怎么设计 Mock 以应对更复杂的实际场景,应该是另外一种艺术了。

2017/12/19 posted in  Python 黑魔法

Werkzeug 常用中间件

昨天记了手脚架 Werkzeug 的几个小工具( Werkzeug 的功能函数 ),今天再补几个中间件。之前在 HackerProxy (原gateway)中,便是通过自定义了一堆 Middleware,来把用户验证等常用逻辑完成的。第一次看到这种俄罗斯套娃式的代码我是震惊的,不过现在认为 Middleware 的确可以适当做一些比较轻的处理。

SharedDataMiddleware

SharedDataMiddleware 是提供一个静态文件分享(下载)的路由。和 flask 中默认的 static 不同,flask 是利用的 send_file ,而 SharedDataMiddleware 可以直接在 app 里注册相关的路由,并绑定一个磁盘路径,并分享这个路径下的文件。

import os
from werkzeug.wsgi import SharedDataMiddleware

app = SharedDataMiddleware(app, {
    '/shared': os.path.join(os.path.dirname(__file__), 'shared')
})

ProfilerMiddleware

ProfilerMiddleware 是一个查看性能的中间件,它会在 profile_dir 下写入访问页面的程序运行状况,包括执行了那些函数,以及运行时间等。

from werkzeug.contrib.profiler import ProfilerMiddleware

app.wsgi_app = ProfilerMiddleware(app, profile_dir="/tmp")

但是如果使用 flask-debugtoolbar 这个拓展,便可以在 Web 上直接查看结果,因此不太建议直接用这个中间件。

DispacherMiddleware

DispatcherMiddleware 是个蛮好玩的中间件,他可以向一个 app 注册其他 app:

from werkzeug.wsgi import DispatcherMiddleware

app.wsgi_app = DispatcherMiddleware(app, {
    '/app2':        app2,
    '/app3':        app3
})

这时,访问 /app2 前缀的 url 就会使用 app2 的相关逻辑。如果对于一个程序可能有分开的需求,可以利用 DispacherMiddleware 做一些插拔操作。

2017/7/13 posted in  Python 黑魔法

Werkzeug 的功能函数

之前在 v2ex 上看到,有人问什么是手脚架,有人举例子,在 Python 中 Werkzeug 便像手脚架,而 flask 就像毛坯房。以此形容 Werkzeug 中有一堆还不错的工具。当然,如果在 flask 中使用 Werkzeug 的话,中间件更是个好东西:Werkzeug 常用中间件

数据结构

Werkzeug 内置里几个常用的数据结构,如果有相关需求就不用重复造轮子了。

TypeConversionDict

如字面意,这是一个类型转换字典,该字典提供一个 get 方法,接受三个参数,key, default, type。

>>> from werkzeug.datastructures import TypeConversionDict
>>> d = TypeConversionDict(foo='42', bar='blub')
>>> d.get('foo', type=int)
42
>>> d.get('bar', -1, type=int)
-1

ImmutableTypeConversionDict

ImmutableTypeConversionDict,便是 TypeConversionDict 的不可变版。

>>> from werkzeug.datastructures import ImmutableTypeConversionDict
>>> dic = ImmutableTypeConversionDict(key="value")
>>> dic.get("key")
'value'
>>> dic["key"] = "v2"
Traceback (most recent call last):
  File "<input>", line 1, in <module>
    dic["k2"] = "v2"
  File "/Library/Python/2.7/site-packages/werkzeug/datastructures.py", line 182, in __setitem__
    is_immutable(self)
  File "/Library/Python/2.7/site-packages/werkzeug/datastructures.py", line 28, in is_immutable
    raise TypeError('%r objects are immutable' % self.__class__.__name__)
TypeError: 'ImmutableTypeConversionDict' objects are immutable

MultiDict

MultiDict 这个东西在我第一次看到时候并没有理解这是做什么的,这是一个允许一个 key 有多个 value 的字典,但是如果当做普通字典用时,只会返回第一个值。

>>> d = MultiDict([('a', 'b'), ('a', 'c')])
>>> d
MultiDict([('a', 'b'), ('a', 'c')])
>>> d['a']
'b'
>>> d.getlist('a')
['b', 'c']
>>> 'a' in d
True

ImmutableMultiDict

ImmutableMultiDict 就是 MultiDict 的不可变版,在 flask 中的 request.args 便是一个 ImmutableMultiDict。

OrderedMultiDict

如字面意,OrderedMultiDict 是 MultiDict 中的 key 按字典序的版。

ImmutableOrderedMultiDict

同上,ImmutableOrderedMultiDict 为 OrderedMultiDict 不可变版。

功能函数

werkzeug.utils 下有一堆好玩的工具,可以在使用 flask 的时候直接拿来用。

cached_property

werkzeug.utils 中有个 cached_property 可以作为装饰器的类,这个装饰器和 @property 有相同的效果,不过被装饰的函数只会第一次运行,然后后面只会返回被缓存的结果。

In [17]: class Foo(object):
    ...:
    ...:     @cached_property
    ...:     def foo(self):
    ...:         print("hello word")
    ...:         return "hello"
    ...:

In [18]: obj = Foo()

In [19]: obj.foo
hello word
Out[19]: 'hello'

In [20]: obj.foo
Out[20]: 'hello'

import_string

import_string 可以通过字符串导出需要导入的模块:

In [21]: from werkzeug.utils import import_string

In [22]: import_string("flask")
Out[22]: <module 'flask' from '/Library/Frameworks/Python.framework/Versions/3.6/lib/python3.6/site-packages/flask/__init__.py'>

In [23]: import_string("flask.Flask")
Out[23]: flask.app.Flask

secure_filename

可以用来生成一个合法的文件名:

>>> secure_filename("My cool movie.mov")
'My_cool_movie.mov'
>>> secure_filename("../../../etc/passwd")
'etc_passwd'
>>> secure_filename(u'i contain cool \xfcml\xe4uts.txt')
'i_contain_cool_umlauts.txt'

密码加密

之前在学校写东西,每次遇到密码加密,都是自己加各种盐,然后找个 md5 的库,加上几遍,对于密码的处理一直有 java 的风范。

Werkzeug 里面是有一套加密和密码验证的工具的。在 werkzeug.security 中,有 generate_password_hash, check_password_hash, 一对好用的工具。

generate_password_hash 接受三个参数:明文密码,加密方法(method, 默认为’pbkdf2:sha256’),盐的长度(salt_length,默认为8)。

In [25]: generate_password_hash("password")
Out[25]: 'pbkdf2:sha256:50000$F6gTN2Eh$52e209ed4431f9268d1bf16295439c46b25ab306acff72615837bdf268fee361'

In [26]: generate_password_hash("password")
Out[26]: 'pbkdf2:sha256:50000$GWo5sU55$729f84eb83e02549312fc4fc51db0614cbc783710902006dc510ccdac0a2b937'

返回的密文是:method$salt$hash 格式。并且,因为有盐的存在,所以同一个密码,并不会有相同的结果。

对于密码的验证:

In [27]: p = generate_password_hash("password")

In [29]: check_password_hash(p, "password")
Out[29]: True

In [30]: check_password_hash(p, "password1")
Out[30]: False
2017/7/12 posted in  Python 黑魔法

Python 生成验证码

在 InfoSub 接入验证码的时候,第一次尝试使用验证码,设计的比较简单,利用的是 captchaitsdangerous 模块。

在创建验证码的时候,后台随机生成一个字符串作为验证码(captcha_code),然后利用 itsdangerous.URLSafeTimedSerializer 对这个字符串进行编码,生成一个安全的(不能被随意反向编码,反向编码需要秘钥),方便的(可以直接扔到 url 里)的字符串。把这个字符串作为验证码的 id(captcha_id)。

这样使得校验验证码非常简单,利用秘钥对 captcha_id 反向编码便可以得到 captcha_code,将 captcha_code 与用户的输入进行比对便可以知道用户的验证码是否正确。

配置

首先要配置用于编码的 serializer:

serializer = URLSafeTimedSerializer(SECRET_KEY)

这样便可以使用 serializerloadsdumps 用于编码和反向编码:

captcha_id = serializer.dumps(captcha_code)
captcha_code = serializer.loads(captcha_id, max_age=max_age)

serializer.loads 方法接受一个 max_age 参数表示过期时间,如果时间超时,将返回空。

生成图片

生成图片用的是 captcha 库,首先是配置:

image = ImageCaptcha(width=160, height=60, fonts=[os.path.join(FONTS_PATH, 'CourierNew-Bold.ttf'),
                                                  os.path.join(FONTS_PATH, 'LiberationMono-Bold.ttf')])

可以指定图片大小,字体,字体大小。

返回图片有两种方式,一个是可以直接拿到图片的 IO 流(一个字节数组),或者直接持久化到硬盘上:

# 获得字节数组
data = image.generate(self.captcha_code)

# 将图片存储到 file_path 这个位置
file_path = tempfile.mkstemp(prefix="INFO-SUB-CAPTCHA-", suffix=".png")
image.write(self.captcha_code, file_path)

考虑到验证码的时效性,便没有采用持久化硬盘上这个操作,而是直接把图片流返回了。

封装的 Captcha 类:

class Captcha(object):
    def __init__(self, captcha_code=None):
        if captcha_code:
            self.captcha_code = captcha_code
        else:
            self.captcha_code = ''.join(random.sample(string.digits + string.lowercase + string.uppercase, 6))
        self.captcha_id = serializer.dumps(self.captcha_code)

    def image(self):
        data = image.generate(self.captcha_code)
        return data

    def save(self):
        file_path = tempfile.mkstemp(prefix="INFO-SUB-CAPTCHA-", suffix=".png")
        image.write(self.captcha_code, file_path)
        return file_path

    def validate(self, code):
        code = code.strip()
        if self.captcha_code.lower() == code.lower():
            return True
        return False

    @classmethod
    def get_by_captcha_id(cls, captcha_id, max_age=60 * 30):
        try:
            captcha_code = serializer.loads(captcha_id, max_age=max_age)
        except:
            captcha_code = None
        if not captcha_code:
            return None
        return cls(captcha_code)

获得验证码图片

因为使用的 flask,而 flask 有 send_file 函数将图片等文件返回:

@view_blueprint.route("/captcha/<captcha_id>")
def get_captcha(captcha_id):
    captcha = Captcha.get_by_captcha_id(captcha_id)
    if not captcha:
        abort(404)
    return send_file(captcha.image(), mimetype='image/png')

这样,直接访问带 captcha_id 的 url 将直接返回验证码图片。

<img src="{{ url_for('view.get_captcha', captcha_id=captcha_id) }}" width="100%" class="img-rounded">

验证操作

为了简单起见,我将验证操作写在了 WTF 表单中,比如登录表单:

class LoginForm(FlaskForm):
    username_or_email = StringField(u"用户名或注册邮箱", validators=[DataRequired(u"登录名不能为空")])
    password = PasswordField(u"密码", validators=[DataRequired(u"密码不能为空")])
    remember = BooleanField(u"记住我", default=False)
    captcha_code = StringField(u"验证码")
    captcha_id = HiddenField()

    def validate_password(self, field):
        user = get_user_by_username_or_email(self.username_or_email.data)
        if user and user.is_active and user.check_password(field.data):
            return
        raise ValidationError(u"用户名或密码错误")

    def validate_captcha_code(self, field):
        if current_app.config.get("DEBUG"):
            return
        if self.captcha_id.data:
            captcha = Captcha.get_by_captcha_id(self.captcha_id.data)
            if captcha.validate(field.data):
                return
        raise ValidationError(u"验证码错误")

在 view 层创建了 form 表单,在返回之前将 captcha_id 填充进去,在提交的时候,便会对验证码进行验证。

2017/7/10 posted in  Python 黑魔法

Python 主机状态监控模块:psutil

这几天毕设需要加一个主机监控的功能,然后发现了一个 python 的跨平台的主机监控模块:psutil。

文档: https://pythonhosted.org/psutil/

Github: https://github.com/giampaolo/psutil

安装

pip install psutil

常用功能

因为毕设中只用到了 CPU,内存,磁盘这三个属性,但是这个库能获得的状态不止这些,还有网络,传感器和各种系统属性。只写下自己用到的,其他的文档介绍的挺全面。

CPU

psutil.cpu_times(percpu=False)

cpu_times 函数接受一个参数,percpu,默认为 False。

>>> psutil.cpu_times()
scputimes(user=33335.82, nice=0.0, system=35253.52, idle=405130.12)

执行之后将返回 user mode 和 kernel mode 的 CPU 时间(user,system),以及 CPU 空转的时间(idle)。

如果加上 percpu 参数,之后,将返回每个核心的信息:

>>> psutil.cpu_times(percpu=True)
[scputimes(user=12113.43, nice=0.0, system=14012.23, idle=92520.62), scputimes(user=4881.28, nice=0.0, system=4959.2, idle=108794.75), scputimes(user=11446.36, nice=0.0, system=11263.96, idle=95924.97), scputimes(user=4969.93, nice=0.0, system=5093.61, idle=108571.62)]

如果只是计算 CPU 的利用率的话,psutil.cpu_percent() 就够解决问题,cpu_percent 可以传递一个间隔参数,来计算一定间隔内的 CPU 利用率:

>>> psutil.cpu_percent(1)
13.7

内存

因为我只需要计算内存的使用率,所以只用到了 virtual_memory 方法。这个方法返回的参数很多,不过也是一目了然。

>>> psutil.virtual_memory()
svmem(total=17179869184L, available=4261830656L, percent=75.2, used=13366374400L, free=2298265600L, active=9087324160L, inactive=1963565056L, wired=2315485184L)

磁盘

disk_partitions 方法可以查看磁盘的使用情况,会返回一个 list,包含了比较全面的信息:

>>> psutil.disk_partitions()
[sdiskpart(device='/dev/disk1', mountpoint='/', fstype='hfs', opts='rw,local,rootfs,dovolfs,journaled,multilabel'), sdiskpart(device='/dev/disk0s3', mountpoint='/Volumes/Recovery HD', fstype='hfs', opts='rw,local,dovolfs,dontbrowse,journaled,multilabel')]

当然,以上只是提到了最简单的功能,还有很多使用的函数可以选择。

2017/6/5 posted in  Python 黑魔法