extension stop hook tests

jupyter_server/extension/manager.py

@@ -297,7 +297,7 @@ def sorted_extensions(self):
     extension_apps = Dict(
         help="""
         Dictionary with extension names as keys
-        and ExtensionApp objects as values.
+        and sets of ExtensionApp objects as values.
         """
     )
 
@@ -352,15 +352,21 @@ def link_extension(self, name, serverapp):
 
     def load_extension(self, name, serverapp):
         extension = self.extensions.get(name)
+
         if extension.enabled:
             try:
-                self.extension_apps.setdefault(name, []).extend(
-                    extension.load_all_points(serverapp)
-                )
-                self.log.info("{name} | extension was successfully loaded.".format(name=name))
+                points = extension.load_all_points(serverapp)
             except Exception as e:
                 self.log.debug("".join(traceback.format_exception(*sys.exc_info())))
                 self.log.warning("{name} | extension failed loading with message: {error}".format(name=name,error=str(e)))
+            else:
+                self.extension_apps.setdefault(name, set()).update((
+                    point
+                    for point in points
+                    if point is not None
+                ))
+                self.log.info("{name} | extension was successfully loaded.".format(name=name))
+
 
     async def stop_extension(self, name, apps):
         """Call the shutdown hooks in the specified apps."""

jupyter_server/tests/extension/test_app.py

@@ -101,3 +101,42 @@ def test_load_parallel_extensions(monkeypatch, jp_environ):
     exts = serverapp.jpserver_extensions
     assert exts['jupyter_server.tests.extension.mockextensions.mock1']
     assert exts['jupyter_server.tests.extension.mockextensions']
+
+
+def test_stop_extension(jp_serverapp, caplog):
+    """Test the stop_extension method.
+
+    This should be fired by ServerApp.cleanup_extensions.
+    """
+    calls = 0
+
+    # load extensions (make sure we only have the one extension loaded
+    jp_serverapp.extension_manager.load_all_extensions(jp_serverapp)
+    assert list(jp_serverapp.extension_manager.extension_apps) == [
+        'jupyter_server.tests.extension.mockextensions'
+    ]
+
+    # add a stop_extension method for the extension app
+    async def _stop(*args):
+        nonlocal calls
+        calls += 1
+    for apps in jp_serverapp.extension_manager.extension_apps.values():
+        for app in apps:
+            if app:
+                app.stop_extension = _stop
+
+    # call cleanup_extensions, check the logging is correct
+    caplog.clear()
+    jp_serverapp.cleanup_extensions()
+    assert [
+        msg
+        for *_, msg in caplog.record_tuples
+    ] == [
+        'Shutting down 1 extension'
+    ]
+
+    # check the extension_apps dictionary is updated
+    assert list(jp_serverapp.extension_manager.extension_apps) == []
+
+    # check the shutdown method was called once
+    assert calls == 1