From 7637dacb3b30d550c6eabd6017a55e164c65c655 Mon Sep 17 00:00:00 2001
From: Oliver Sanders <oliver.sanders@metoffice.gov.uk>
Date: Fri, 11 Jun 2021 16:40:22 +0100
Subject: [PATCH] extension stop hook tests

---
 jupyter_server/extension/manager.py        | 16 ++++++---
 jupyter_server/tests/extension/test_app.py | 39 ++++++++++++++++++++++
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/jupyter_server/extension/manager.py b/jupyter_server/extension/manager.py
index 42c24196da..62d9c95c48 100644
--- a/jupyter_server/extension/manager.py
+++ b/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."""
diff --git a/jupyter_server/tests/extension/test_app.py b/jupyter_server/tests/extension/test_app.py
index 3cc0e82fe6..25e54cfcfc 100644
--- a/jupyter_server/tests/extension/test_app.py
+++ b/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