Merge pull request #1557 from ssanderson/addcleanup-instead-of-teardown

Use addCleanup instead of tearDown in several test suites.

notebook/nbconvert/tests/test_nbconvert_handlers.py

@@ -58,8 +58,17 @@ class APITest(NotebookTestBase):
         nbdir = self.notebook_dir.name
         
         if not os.path.isdir(pjoin(nbdir, 'foo')):
-            os.mkdir(pjoin(nbdir, 'foo'))
-        
+            subdir = pjoin(nbdir, 'foo')
+
+            os.mkdir(subdir)
+
+            # Make sure that we clean this up when we're done.
+            # By using addCleanup this will happen correctly even if we fail
+            # later in setUp.
+            @self.addCleanup
+            def cleanup_dir():
+                shutil.rmtree(subdir, ignore_errors=True)
+
         nb = new_notebook()
         
         nb.cells.append(new_markdown_cell(u'Created by test ³'))
@@ -77,12 +86,6 @@ class APITest(NotebookTestBase):
 
         self.nbconvert_api = NbconvertAPI(self.base_url())
 
-    def tearDown(self):
-        nbdir = self.notebook_dir.name
-
-        for dname in ['foo']:
-            shutil.rmtree(pjoin(nbdir, dname), ignore_errors=True)
-    
     @onlyif_cmds_exist('pandoc')
     def test_from_file(self):
         r = self.nbconvert_api.from_file('html', 'foo', 'testnb.ipynb')

notebook/services/contents/tests/test_contents_api.py

@@ -2,6 +2,7 @@
 """Test the contents webservice API."""
 
 from contextlib import contextmanager
+from functools import partial
 import io
 import json
 import os
@@ -195,32 +196,32 @@ class APITest(NotebookTestBase):
     
     def isdir(self, api_path):
         return os.path.isdir(self.to_os_path(api_path))
-    
-    def setUp(self):
 
+    def setUp(self):
         for d in (self.dirs + self.hidden_dirs):
             self.make_dir(d)
+            self.addCleanup(partial(self.delete_dir, d))
 
         for d, name in self.dirs_nbs:
             # create a notebook
             nb = new_notebook()
-            self.make_nb(u'{}/{}.ipynb'.format(d, name), nb)
-            
+            nbname = u'{}/{}.ipynb'.format(d, name)
+            self.make_nb(nbname, nb)
+            self.addCleanup(partial(self.delete_file, nbname))
+
             # create a text file
             txt = self._txt_for_name(name)
-            self.make_txt(u'{}/{}.txt'.format(d, name), txt)
-            
-            # create a binary file
+            txtname = u'{}/{}.txt'.format(d, name)
+            self.make_txt(txtname, txt)
+            self.addCleanup(partial(self.delete_file, txtname))
+
             blob = self._blob_for_name(name)
-            self.make_blob(u'{}/{}.blob'.format(d, name), blob)
+            blobname = u'{}/{}.blob'.format(d, name)
+            self.make_blob(blobname, blob)
+            self.addCleanup(partial(self.delete_file, blobname))
 
         self.api = API(self.base_url())
 
-    def tearDown(self):
-        for dname in (list(self.top_level_dirs) + self.hidden_dirs):
-            self.delete_dir(dname)
-        self.delete_file('inroot.ipynb')
-
     def test_list_notebooks(self):
         nbs = notebooks_only(self.api.list().json())
         self.assertEqual(len(nbs), 1)

notebook/services/sessions/tests/test_sessionmanager.py

@@ -1,5 +1,6 @@
 """Tests for the session manager."""
 
+from functools import partial
 from unittest import TestCase
 
 from tornado import gen, web
@@ -39,10 +40,8 @@ class TestSessionManager(TestCase):
             contents_manager=ContentsManager(),
         )
         self.loop = IOLoop()
-    
-    def tearDown(self):
-        self.loop.close(all_fds=True)
-    
+        self.addCleanup(partial(self.loop.close, all_fds=True))
+
     def create_sessions(self, *kwarg_list):
         @gen.coroutine
         def co_add():

notebook/services/sessions/tests/test_sessions_api.py

@@ -1,6 +1,7 @@
 """Test the sessions web service API."""
 
 import errno
+from functools import partial
 import io
 import os
 import json
@@ -65,33 +66,34 @@ class SessionAPITest(NotebookTestBase):
     """Test the sessions web service API"""
     def setUp(self):
         nbdir = self.notebook_dir.name
+        subdir = pjoin(nbdir, 'foo')
+
         try:
-            os.mkdir(pjoin(nbdir, 'foo'))
+            os.mkdir(subdir)
         except OSError as e:
             # Deleting the folder in an earlier test may have failed
             if e.errno != errno.EEXIST:
                 raise
+        self.addCleanup(partial(shutil.rmtree, subdir, ignore_errors=True))
 
-        with io.open(pjoin(nbdir, 'foo', 'nb1.ipynb'), 'w',
-                     encoding='utf-8') as f:
+        with io.open(pjoin(subdir, 'nb1.ipynb'), 'w', encoding='utf-8') as f:
             nb = new_notebook()
             write(nb, f, version=4)
 
         self.sess_api = SessionAPI(self.base_url())
 
-    def tearDown(self):
-        for session in self.sess_api.list().json():
-            self.sess_api.delete(session['id'])
-        # This is necessary in some situations on Windows: without it, it
-        # fails to delete the directory because something is still using it. I
-        # think there is a brief period after the kernel terminates where
-        # Windows still treats its working directory as in use. On my Windows
-        # VM, 0.01s is not long enough, but 0.1s appears to work reliably.
-        # -- TK, 15 December 2014
-        time.sleep(0.1)
-
-        shutil.rmtree(pjoin(self.notebook_dir.name, 'foo'),
-                      ignore_errors=True)
+        @self.addCleanup
+        def cleanup_sessions():
+            for session in self.sess_api.list().json():
+                self.sess_api.delete(session['id'])
+
+            # This is necessary in some situations on Windows: without it, it
+            # fails to delete the directory because something is still using
+            # it. I think there is a brief period after the kernel terminates
+            # where Windows still treats its working directory as in use. On my
+            # Windows VM, 0.01s is not long enough, but 0.1s appears to work
+            # reliably.  -- TK, 15 December 2014
+            time.sleep(0.1)
 
     def test_create(self):
         sessions = self.sess_api.list().json()

notebook/tests/test_nbextensions.py

@@ -62,7 +62,15 @@ class TestInstallNBExtension(TestCase):
         return py3compat.cast_unicode(td.name)
 
     def setUp(self):
+        # Any TemporaryDirectory objects appended to this list will be cleaned
+        # up at the end of the test run.
         self.tempdirs = []
+
+        @self.addCleanup
+        def cleanup_tempdirs():
+            for d in self.tempdirs:
+                d.cleanup()
+
         self.src = self.tempdir()
         self.files = files = [
             pjoin(u'ƒile'),
@@ -75,25 +83,30 @@ class TestInstallNBExtension(TestCase):
             if not os.path.exists(parent):
                 os.makedirs(parent)
             touch(fullpath)
-        
+
         self.test_dir = self.tempdir()
         self.data_dir = os.path.join(self.test_dir, 'data')
         self.config_dir = os.path.join(self.test_dir, 'config')
         self.system_data_dir = os.path.join(self.test_dir, 'system_data')
         self.system_path = [self.system_data_dir]
         self.system_nbext = os.path.join(self.system_data_dir, 'nbextensions')
+
+        # Patch out os.environ so that tests are isolated from the real OS
+        # environment.
         self.patch_env = patch.dict('os.environ', {
             'JUPYTER_CONFIG_DIR': self.config_dir,
             'JUPYTER_DATA_DIR': self.data_dir,
         })
         self.patch_env.start()
-        self.patch_system_path = patch.object(nbextensions,
-            'SYSTEM_JUPYTER_PATH', self.system_path)
+        self.addCleanup(self.patch_env.stop)
+
+        # Patch out the system path os that we consistently use our own
+        # temporary directory instead.
+        self.patch_system_path = patch.object(
+            nbextensions, 'SYSTEM_JUPYTER_PATH', self.system_path
+        )
         self.patch_system_path.start()
-    
-    def tearDown(self):
-        self.patch_env.stop()
-        self.patch_system_path.stop()
+        self.addCleanup(self.patch_system_path.stop)
 
     def assert_dir_exists(self, path):
         if not os.path.exists(path):