Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

How can I unit test a class with __enter__ / __exit__ methods?

I am new to unit testing and I am trying to find a way to test whether the with keyword is working correctly in my object.

In this case, my object has an __enter__ method which creates a temporary directory and __exit__ method which is supposed to destroy it. (It also has a do_stuff method that I only included to test writing to the temporary directory.)

I'm not entirely sure how to approach testing this. I've checked out the unittest module, and even written some tests for basic methods before, but I am not sure what the best way is in this case... or if this even makes sense. Anyway, here is my code for the object:

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()
like image 605
KaleidoEscape Avatar asked Oct 20 '22 23:10

KaleidoEscape


1 Answers

If you want to test that the MyObj works with with statement, and it create/delete temporary directory, use the with statement in the test method:

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()
like image 141
falsetru Avatar answered Oct 31 '22 18:10

falsetru