首页 > 解决方案 > 在Unittest中用setUp和tearDown替换python with_statement

问题描述

在一个测试套件中,我有一些代码组织如下,上下文是一些在退出with块时被删除的持久对象:

class Test(TestCase):
    def test_action1(self):
        with create_context() as context:
            context.prepare_context()
            context.action1()
            self.assertTrue(context.check1())
            
    def test_action2(self):
        with create_context() as context:
            context.prepare_context()
            context.action2()
            self.assertTrue(context.check2())

很明显,代码在两个测试中都有一些重复的设置样板,因此我想使用 setUp() 和 tearDown() 方法来分解该样板。

但我不知道如何提取 with_statement。我想出的是这样的:

class Test(TestCase):
    def setUp(self):
        self.context = create_context()
        self.context.prepare_context()

    def tearDown(self):
        del self.context()

    def test_action1(self):
        self.context.action1()
        self.assertTrue(self.context.check1())
            
    def test_action2(self):
        self.context.action2()
        self.assertTrue(self.context.check2())

但是我相信当测试失败时这并不等同,而且必须在 tearDown() 中明确删除也感觉不对。

将我的 with_statement 代码更改为 setUp() 和 tearDown() 样式的正确方法是什么?

标签: pythonpython-unittestwith-statement

解决方案


您可能想使用contextlib.ExitStack它。

一种上下文管理器,旨在使以编程方式组合其他上下文管理器和清理功能变得容易,尤其是那些可选的或由输入数据驱动的那些。

import contextlib
from unittest import TestCase


class Test(TestCase):
    def setUp(self) -> None:
        stack = contextlib.ExitStack()
        self.context = stack.enter_context(create_context())  # create_context is your context manager
        self.addCleanup(stack.close)

    def test_action1(self):
        self.context.prepare_context()
        self.context.action1()
        self.assertTrue(self.context.check1())

或者如果您想对拆解进行一些控制或使用多个上下文管理器,这会更好

import contextlib
from unittest import TestCase


class Test(TestCase):
    def setUp(self):
        with contextlib.ExitStack() as stack:
            self.context = stack.enter_context(create_context())  # create_context is your context manager
            self._resource_stack = stack.pop_all()

    def tearDown(self):
        self._resource_stack.close()

    def test_action1(self):
        self.context.prepare_context()
        self.context.action1()
        self.assertTrue(self.context.check1())

推荐阅读