python - 如何使用 __enter__/__exit__ 方法对类进行单元测试？

Question

我是单元测试的新手，我正在尝试找到一种方法来测试 with 关键字是否在我的对象中正常工作。

在这种情况下，我的对象有一个创建临时目录的 __enter__ 方法和应该销毁它的 __exit__ 方法。 (它还有一个 do_stuff 方法，我只包含它来测试写入临时目录。)

我不完全确定如何进行测试。我检查了 unittest 模块，之前甚至为基本方法编写了一些测试，但我不确定在这种情况下最好的方法是什么……或者这是否有意义。无论如何，这是我的对象代码:

import shutil
import tempfile
import os
import glob

class MyObj(object):
    def __enter__(self):
        self.tmpdir = tempfile.mkdtemp(dir='.')
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        shutil.rmtree(self.tmpdir)

    def do_stuff(self):
        new = os.path.join(self.tmpdir, 'new_file.txt')
        with open(new, 'w') as nf:
            nf.write('testing')
        print(glob.glob(self.tmpdir + '/*'))

myobj = MyObj()
with myobj as x:
    x.do_stuff()

Answer 1

如果你想测试 MyObj 是否与 with 语句一起工作，并且它创建/删除临时目录，请使用 with 语句测试方法:

import unittest


class TestMyObj(unittest.TestCase):

    def test_myobj_with_statement__should_create_delete_temp_directory(self):
        with MyObj() as obj:
            # Directory is created
            self.assertTrue(os.path.isdir(obj.tmpdir))
        # Directory is gone
        self.assertFalse(os.path.isdir(obj.tmpdir))

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