check for leftover zmq resources across tests

Author: minrk

jupyter_client/tests/conftest.py

@@ -1,9 +1,12 @@
 import asyncio
+import gc
 import os
 import sys
 
 import pytest
+import zmq
 from jupyter_core import paths
+from zmq.tests import BaseZMQTestCase
 
 from .utils import test_env
 
@@ -39,3 +42,47 @@ def env():
 @pytest.fixture()
 def kernel_dir():
     return pjoin(paths.jupyter_data_dir(), 'kernels')
+
+
+def assert_no_zmq():
+    """Verify that there are no zmq resources
+
+    avoids reference leaks across tests,
+    which can lead to FD exhaustion
+    """
+    # zmq garbage collection uses a zmq socket in a thread
+    # we don't want to delete these from the main thread!
+    from zmq.utils import garbage
+
+    garbage.gc.stop()
+    sockets = [
+        obj
+        for obj in gc.get_referrers(zmq.Socket)
+        if isinstance(obj, zmq.Socket) and not obj.closed
+    ]
+    if sockets:
+        message = f"{len(sockets)} unclosed sockets: {sockets}"
+        for s in sockets:
+            s.close(linger=0)
+        raise AssertionError(message)
+    contexts = [
+        obj
+        for obj in gc.get_referrers(zmq.Context)
+        if isinstance(obj, zmq.Context) and not obj.closed
+    ]
+    # allow for single zmq.Context.instance()
+    if contexts and len(contexts) > 1:
+        message = f"{len(contexts)} unclosed contexts: {contexts}"
+        for ctx in contexts:
+            ctx.destroy(linger=0)
+        raise AssertionError(message)
+
+
+@pytest.fixture(autouse=True)
+def check_zmq(request):
+    yield
+    if request.instance and isinstance(request.instance, BaseZMQTestCase):
+        # can't run this check on old-style TestCases with tearDown methods
+        # because this check runs before tearDown
+        return
+    assert_no_zmq()

jupyter_client/tests/test_multikernelmanager.py

@@ -69,6 +69,7 @@ def _run_lifecycle(km, test_kid=None):
         assert isinstance(k, KernelManager)
         km.shutdown_kernel(kid, now=True)
         assert kid not in km, f"{kid} not in {km}"
+        km.context.term()
 
     def _run_cinfo(self, km, transport, ip):
         kid = km.start_kernel(stdout=PIPE, stderr=PIPE)
@@ -87,6 +88,7 @@ def _run_cinfo(self, km, transport, ip):
         stream = km.connect_hb(kid)
         stream.close()
         km.shutdown_kernel(kid, now=True)
+        km.context.term()
 
     # static so picklable for multiprocessing on Windows
     @classmethod
@@ -106,6 +108,7 @@ def test_shutdown_all(self):
         self.assertNotIn(kid, km)
         # shutdown again is okay, because we have no kernels
         km.shutdown_all()
+        km.context.term()
 
     def test_tcp_cinfo(self):
         km = self._get_tcp_km()
@@ -217,6 +220,7 @@ def test_subclass_callables(self):
         assert km.call_count("cleanup_resources") == 0
 
         assert kid not in km, f"{kid} not in {km}"
+        km.context.term()
 
 
 class TestAsyncKernelManager(AsyncTestCase):
@@ -263,6 +267,7 @@ async def _run_lifecycle(km, test_kid=None):
         assert isinstance(k, AsyncKernelManager)
         await km.shutdown_kernel(kid, now=True)
         assert kid not in km, f"{kid} not in {km}"
+        km.context.term()
 
     async def _run_cinfo(self, km, transport, ip):
         kid = await km.start_kernel(stdout=PIPE, stderr=PIPE)
@@ -282,6 +287,7 @@ async def _run_cinfo(self, km, transport, ip):
         stream.close()
         await km.shutdown_kernel(kid, now=True)
         self.assertNotIn(kid, km)
+        km.context.term()
 
     @gen_test
     async def test_tcp_lifecycle(self):
@@ -316,6 +322,7 @@ async def test_use_after_shutdown_all(self):
         self.assertNotIn(kid, km)
         # shutdown again is okay, because we have no kernels
         await km.shutdown_all()
+        km.context.term()
 
     @gen_test(timeout=20)
     async def test_shutdown_all_while_starting(self):
@@ -333,6 +340,7 @@ async def test_shutdown_all_while_starting(self):
         self.assertNotIn(kid, km)
         # shutdown again is okay, because we have no kernels
         await km.shutdown_all()
+        km.context.term()
 
     @gen_test
     async def test_tcp_cinfo(self):
@@ -466,3 +474,4 @@ async def test_subclass_callables(self):
         assert mkm.call_count("cleanup_resources") == 0
 
         assert kid not in mkm, f"{kid} not in {mkm}"
+        mkm.context.term()