From 8a0b3437600aebf08f91e886cb272dc627bef3b8 Mon Sep 17 00:00:00 2001
From: David Wilson <dw@botanicus.net>
Date: Sat, 3 Nov 2018 13:28:37 +0000
Subject: [PATCH] issue #406: test for FD leak after every TestCase

---
 dev_requirements.txt        |  1 +
 tests/call_error_test.py    |  4 ++--
 tests/docker_test.py        |  2 +-
 tests/fakessh_test.py       |  2 +-
 tests/fork_test.py          |  5 +++--
 tests/local_test.py         |  4 ++--
 tests/master_test.py        |  2 +-
 tests/minify_test.py        |  4 ++--
 tests/parent_test.py        |  8 ++++----
 tests/responder_test.py     |  8 ++++----
 tests/router_test.py        |  7 +++----
 tests/serialization_test.py |  4 ++--
 tests/ssh_test.py           |  4 ++--
 tests/testlib.py            | 26 ++++++++++++++++++++++++++
 tests/types_test.py         |  6 ++++--
 tests/unix_test.py          |  6 +++---
 tests/utils_test.py         |  8 +++++---
 17 files changed, 66 insertions(+), 35 deletions(-)

diff --git a/dev_requirements.txt b/dev_requirements.txt
index f48006e5..c536c154 100644
--- a/dev_requirements.txt
+++ b/dev_requirements.txt
@@ -1,4 +1,5 @@
 -r docs/docs-requirements.txt
+psutil==5.4.8
 coverage==4.5.1
 Django==1.6.11  # Last version supporting 2.6.
 mock==2.0.0
diff --git a/tests/call_error_test.py b/tests/call_error_test.py
index 447a80a9..1480a743 100644
--- a/tests/call_error_test.py
+++ b/tests/call_error_test.py
@@ -10,7 +10,7 @@ import testlib
 import plain_old_module
 
 
-class ConstructorTest(unittest2.TestCase):
+class ConstructorTest(testlib.TestCase):
     klass = mitogen.core.CallError
 
     def test_string_noargs(self):
@@ -44,7 +44,7 @@ class ConstructorTest(unittest2.TestCase):
         self.assertTrue('test_from_exc_tb' in e.args[0])
 
 
-class PickleTest(unittest2.TestCase):
+class PickleTest(testlib.TestCase):
     klass = mitogen.core.CallError
 
     def test_string_noargs(self):
diff --git a/tests/docker_test.py b/tests/docker_test.py
index 2d45609a..49c742ee 100644
--- a/tests/docker_test.py
+++ b/tests/docker_test.py
@@ -7,7 +7,7 @@ import unittest2
 import testlib
 
 
-class ConstructorTest(testlib.RouterMixin, unittest2.TestCase):
+class ConstructorTest(testlib.RouterMixin, testlib.TestCase):
     def test_okay(self):
         docker_path = testlib.data_path('stubs/stub-docker.py')
         context = self.router.docker(
diff --git a/tests/fakessh_test.py b/tests/fakessh_test.py
index c584acfe..63c70058 100644
--- a/tests/fakessh_test.py
+++ b/tests/fakessh_test.py
@@ -10,7 +10,7 @@ import mitogen.fakessh
 import testlib
 
 
-class RsyncTest(testlib.DockerMixin, unittest2.TestCase):
+class RsyncTest(testlib.DockerMixin, testlib.TestCase):
     @timeoutcontext.timeout(5)
     @unittest2.skip('broken')
     def test_rsync_from_master(self):
diff --git a/tests/fork_test.py b/tests/fork_test.py
index 8b396bbf..5e457c97 100644
--- a/tests/fork_test.py
+++ b/tests/fork_test.py
@@ -55,7 +55,7 @@ def exercise_importer(n):
     return simple_pkg.a.subtract_one_add_two(n)
 
 
-class ForkTest(testlib.RouterMixin, unittest2.TestCase):
+class ForkTest(testlib.RouterMixin, testlib.TestCase):
     def test_okay(self):
         context = self.router.fork()
         self.assertNotEqual(context.call(os.getpid), os.getpid())
@@ -84,7 +84,8 @@ class ForkTest(testlib.RouterMixin, unittest2.TestCase):
         context = self.router.fork(on_start=on_start)
         self.assertEquals(123, recv.get().unpickle())
 
-class DoubleChildTest(testlib.RouterMixin, unittest2.TestCase):
+
+class DoubleChildTest(testlib.RouterMixin, testlib.TestCase):
     def test_okay(self):
         # When forking from the master process, Mitogen had nothing to do with
         # setting up stdio -- that was inherited wherever the Master is running
diff --git a/tests/local_test.py b/tests/local_test.py
index fbf5c1c8..5a620e52 100644
--- a/tests/local_test.py
+++ b/tests/local_test.py
@@ -20,7 +20,7 @@ def get_os_environ():
     return dict(os.environ)
 
 
-class LocalTest(testlib.RouterMixin, unittest2.TestCase):
+class LocalTest(testlib.RouterMixin, testlib.TestCase):
     stream_class = mitogen.ssh.Stream
 
     def test_stream_name(self):
@@ -29,7 +29,7 @@ class LocalTest(testlib.RouterMixin, unittest2.TestCase):
         self.assertEquals('local.%d' % (pid,), context.name)
 
 
-class PythonPathTest(testlib.RouterMixin, unittest2.TestCase):
+class PythonPathTest(testlib.RouterMixin, testlib.TestCase):
     stream_class = mitogen.ssh.Stream
 
     def test_inherited(self):
diff --git a/tests/master_test.py b/tests/master_test.py
index 19a9b414..31d11013 100644
--- a/tests/master_test.py
+++ b/tests/master_test.py
@@ -6,7 +6,7 @@ import testlib
 import mitogen.master
 
 
-class ScanCodeImportsTest(unittest2.TestCase):
+class ScanCodeImportsTest(testlib.TestCase):
     func = staticmethod(mitogen.master.scan_code_imports)
 
     if mitogen.core.PY3:
diff --git a/tests/minify_test.py b/tests/minify_test.py
index 98307059..e990fb90 100644
--- a/tests/minify_test.py
+++ b/tests/minify_test.py
@@ -16,7 +16,7 @@ def read_sample(fname):
     return sample
 
 
-class MinimizeSourceTest(unittest2.TestCase):
+class MinimizeSourceTest(testlib.TestCase):
     func = staticmethod(mitogen.minify.minimize_source)
 
     def test_class(self):
@@ -55,7 +55,7 @@ class MinimizeSourceTest(unittest2.TestCase):
         self.assertEqual(expected, self.func(original))
 
 
-class MitogenCoreTest(unittest2.TestCase):
+class MitogenCoreTest(testlib.TestCase):
     # Verify minimize_source() succeeds for all built-in modules.
     func = staticmethod(mitogen.minify.minimize_source)
 
diff --git a/tests/parent_test.py b/tests/parent_test.py
index 9d540ccc..e6a93deb 100644
--- a/tests/parent_test.py
+++ b/tests/parent_test.py
@@ -153,7 +153,7 @@ class StreamErrorTest(testlib.RouterMixin, testlib.TestCase):
         self.assertTrue(s in e.args[0])
 
 
-class ContextTest(testlib.RouterMixin, unittest2.TestCase):
+class ContextTest(testlib.RouterMixin, testlib.TestCase):
     def test_context_shutdown(self):
         local = self.router.local()
         pid = local.call(os.getpid)
@@ -181,7 +181,7 @@ class OpenPtyTest(testlib.TestCase):
         self.assertEquals(e.args[0], msg)
 
 
-class TtyCreateChildTest(unittest2.TestCase):
+class TtyCreateChildTest(testlib.TestCase):
     func = staticmethod(mitogen.parent.tty_create_child)
 
     def test_dev_tty_open_succeeds(self):
@@ -211,7 +211,7 @@ class TtyCreateChildTest(unittest2.TestCase):
             tf.close()
 
 
-class IterReadTest(unittest2.TestCase):
+class IterReadTest(testlib.TestCase):
     func = staticmethod(mitogen.parent.iter_read)
 
     def make_proc(self):
@@ -263,7 +263,7 @@ class IterReadTest(unittest2.TestCase):
             proc.terminate()
 
 
-class WriteAllTest(unittest2.TestCase):
+class WriteAllTest(testlib.TestCase):
     func = staticmethod(mitogen.parent.write_all)
 
     def make_proc(self):
diff --git a/tests/responder_test.py b/tests/responder_test.py
index 46400fce..888302c0 100644
--- a/tests/responder_test.py
+++ b/tests/responder_test.py
@@ -13,7 +13,7 @@ import plain_old_module
 import simple_pkg.a
 
 
-class NeutralizeMainTest(testlib.RouterMixin, unittest2.TestCase):
+class NeutralizeMainTest(testlib.RouterMixin, testlib.TestCase):
     klass = mitogen.master.ModuleResponder
 
     def call(self, *args, **kwargs):
@@ -67,7 +67,7 @@ class NeutralizeMainTest(testlib.RouterMixin, unittest2.TestCase):
 
 
 
-class GoodModulesTest(testlib.RouterMixin, unittest2.TestCase):
+class GoodModulesTest(testlib.RouterMixin, testlib.TestCase):
     def test_plain_old_module(self):
         # The simplest case: a top-level module with no interesting imports or
         # package machinery damage.
@@ -89,7 +89,7 @@ class GoodModulesTest(testlib.RouterMixin, unittest2.TestCase):
         self.assertEquals(output, "['__main__', 50]\n")
 
 
-class BrokenModulesTest(unittest2.TestCase):
+class BrokenModulesTest(testlib.TestCase):
     def test_obviously_missing(self):
         # Ensure we don't crash in the case of a module legitimately being
         # unavailable. Should never happen in the real world.
@@ -144,7 +144,7 @@ class BrokenModulesTest(unittest2.TestCase):
         self.assertIsInstance(msg.unpickle(), tuple)
 
 
-class BlacklistTest(unittest2.TestCase):
+class BlacklistTest(testlib.TestCase):
     @unittest2.skip('implement me')
     def test_whitelist_no_blacklist(self):
         assert 0
diff --git a/tests/router_test.py b/tests/router_test.py
index 7b7e2896..b0add6d3 100644
--- a/tests/router_test.py
+++ b/tests/router_test.py
@@ -36,7 +36,7 @@ def send_n_sized_reply(sender, n):
     return 123
 
 
-class SourceVerifyTest(testlib.RouterMixin, unittest2.TestCase):
+class SourceVerifyTest(testlib.RouterMixin, testlib.TestCase):
     def setUp(self):
         super(SourceVerifyTest, self).setUp()
         # Create some children, ping them, and store what their messages look
@@ -149,7 +149,7 @@ class PolicyTest(testlib.RouterMixin, testlib.TestCase):
         self.assertEquals(e.args[0], self.router.refused_msg)
 
 
-class CrashTest(testlib.BrokerMixin, unittest2.TestCase):
+class CrashTest(testlib.BrokerMixin, testlib.TestCase):
     # This is testing both Broker's ability to crash nicely, and Router's
     # ability to respond to the crash event.
     klass = mitogen.master.Router
@@ -178,8 +178,7 @@ class CrashTest(testlib.BrokerMixin, unittest2.TestCase):
         self.assertTrue(expect in log.stop())
 
 
-
-class AddHandlerTest(unittest2.TestCase):
+class AddHandlerTest(testlib.TestCase):
     klass = mitogen.master.Router
 
     def test_invoked_at_shutdown(self):
diff --git a/tests/serialization_test.py b/tests/serialization_test.py
index f108ff37..d8c54c59 100644
--- a/tests/serialization_test.py
+++ b/tests/serialization_test.py
@@ -20,7 +20,7 @@ def roundtrip(v):
     return mitogen.core.Message(data=msg.data).unpickle()
 
 
-class BlobTest(unittest2.TestCase):
+class BlobTest(testlib.TestCase):
     klass = mitogen.core.Blob
 
     # Python 3 pickle protocol 2 does weird stuff depending on whether an empty
@@ -36,7 +36,7 @@ class BlobTest(unittest2.TestCase):
         self.assertEquals(b(''), roundtrip(v))
 
 
-class ContextTest(testlib.RouterMixin, unittest2.TestCase):
+class ContextTest(testlib.RouterMixin, testlib.TestCase):
     klass = mitogen.core.Context
 
     # Ensure Context can be round-tripped by regular pickle in addition to
diff --git a/tests/ssh_test.py b/tests/ssh_test.py
index 36359a66..661ff5ed 100644
--- a/tests/ssh_test.py
+++ b/tests/ssh_test.py
@@ -29,7 +29,7 @@ class StubSshMixin(testlib.RouterMixin):
             del os.environ['STUBSSH_MODE']
 
 
-class ConstructorTest(testlib.RouterMixin, unittest2.TestCase):
+class ConstructorTest(testlib.RouterMixin, testlib.TestCase):
     def test_okay(self):
         context = self.router.ssh(
             hostname='hostname',
@@ -165,7 +165,7 @@ class SshTest(testlib.DockerMixin, testlib.TestCase):
             fp.close()
 
 
-class BannerTest(testlib.DockerMixin, unittest2.TestCase):
+class BannerTest(testlib.DockerMixin, testlib.TestCase):
     # Verify the ability to disambiguate random spam appearing in the SSHd's
     # login banner from a legitimate password prompt.
     stream_class = mitogen.ssh.Stream
diff --git a/tests/testlib.py b/tests/testlib.py
index 8f11337d..605254d3 100644
--- a/tests/testlib.py
+++ b/tests/testlib.py
@@ -8,9 +8,11 @@ import subprocess
 import sys
 import time
 
+import psutil
 import unittest2
 
 import mitogen.core
+import mitogen.fork
 import mitogen.master
 import mitogen.utils
 
@@ -41,6 +43,13 @@ if faulthandler is not None:
     faulthandler.enable()
 
 
+def get_fd_count():
+    """
+    Return the number of FDs open by this process.
+    """
+    return psutil.Process().num_fds()
+
+
 def data_path(suffix):
     path = os.path.join(DATA_DIR, suffix)
     if path.endswith('.key'):
@@ -211,6 +220,23 @@ class LogCapturer(object):
 
 
 class TestCase(unittest2.TestCase):
+    @classmethod
+    def setUpClass(cls):
+        # This is done in setUpClass() so we have a chance to run before any
+        # Broker() instantiations in setUp() etc.
+        mitogen.fork.on_fork()
+        cls._fd_count_before = get_fd_count()
+        super(TestCase, cls).setUpClass()
+
+    @classmethod
+    def tearDownClass(cls):
+        super(TestCase, cls).tearDownClass()
+        mitogen.fork.on_fork()
+        assert get_fd_count() == cls._fd_count_before, \
+            "%s leaked FDs. Count before: %s, after: %s" % (
+                cls, cls._fd_count_before, get_fd_count(),
+        )
+
     def assertRaises(self, exc, func, *args, **kwargs):
         """Like regular assertRaises, except return the exception that was
         raised. Can't use context manager because tests must run on Python2.4"""
diff --git a/tests/types_test.py b/tests/types_test.py
index 4f80e076..f929c098 100644
--- a/tests/types_test.py
+++ b/tests/types_test.py
@@ -11,8 +11,10 @@ import unittest2
 import mitogen.core
 from mitogen.core import b
 
+import testlib
 
-class BlobTest(unittest2.TestCase):
+
+class BlobTest(testlib.TestCase):
     klass = mitogen.core.Blob
 
     def make(self):
@@ -43,7 +45,7 @@ class BlobTest(unittest2.TestCase):
                           mitogen.core.BytesType(blob2))
 
 
-class SecretTest(unittest2.TestCase):
+class SecretTest(testlib.TestCase):
     klass = mitogen.core.Secret
 
     def make(self):
diff --git a/tests/unix_test.py b/tests/unix_test.py
index 67265c81..f837c6f0 100644
--- a/tests/unix_test.py
+++ b/tests/unix_test.py
@@ -30,7 +30,7 @@ class MyService(mitogen.service.Service):
         }
 
 
-class IsPathDeadTest(unittest2.TestCase):
+class IsPathDeadTest(testlib.TestCase):
     func = staticmethod(mitogen.unix.is_path_dead)
     path = '/tmp/stale-socket'
 
@@ -57,7 +57,7 @@ class IsPathDeadTest(unittest2.TestCase):
         os.unlink(self.path)
 
 
-class ListenerTest(testlib.RouterMixin, unittest2.TestCase):
+class ListenerTest(testlib.RouterMixin, testlib.TestCase):
     klass = mitogen.unix.Listener
 
     def test_constructor_basic(self):
@@ -66,7 +66,7 @@ class ListenerTest(testlib.RouterMixin, unittest2.TestCase):
         os.unlink(listener.path)
 
 
-class ClientTest(unittest2.TestCase):
+class ClientTest(testlib.TestCase):
     klass = mitogen.unix.Listener
 
     def _try_connect(self, path):
diff --git a/tests/utils_test.py b/tests/utils_test.py
index b2e0aa9e..5b81289e 100644
--- a/tests/utils_test.py
+++ b/tests/utils_test.py
@@ -6,6 +6,8 @@ import mitogen.core
 import mitogen.master
 import mitogen.utils
 
+import testlib
+
 
 def func0(router):
     return router
@@ -16,7 +18,7 @@ def func(router):
     return router
 
 
-class RunWithRouterTest(unittest2.TestCase):
+class RunWithRouterTest(testlib.TestCase):
     # test_shutdown_on_exception
     # test_shutdown_on_success
 
@@ -26,7 +28,7 @@ class RunWithRouterTest(unittest2.TestCase):
         self.assertFalse(router.broker._thread.isAlive())
 
 
-class WithRouterTest(unittest2.TestCase):
+class WithRouterTest(testlib.TestCase):
     def test_with_broker(self):
         router = func()
         self.assertIsInstance(router, mitogen.master.Router)
@@ -40,7 +42,7 @@ class Unicode(mitogen.core.UnicodeType): pass
 class Bytes(mitogen.core.BytesType): pass
 
 
-class CastTest(unittest2.TestCase):
+class CastTest(testlib.TestCase):
     def test_dict(self):
         self.assertEqual(type(mitogen.utils.cast({})), dict)
         self.assertEqual(type(mitogen.utils.cast(Dict())), dict)