From 0d36baa4b000fe1bf67b2c6e33ac13113e34d7fd Mon Sep 17 00:00:00 2001
From: Thomas Wilson <t.wilson.au@gmail.com>
Date: Thu, 21 May 2020 15:20:43 +0800
Subject: [PATCH] Add better plugin unloading

---
 shepherd/agent/plugin.py | 47 +++++++++++++++++++++++++++++++++++++++-
 tests/test_core.py       |  2 ++
 tests/test_plugins.py    | 16 ++++----------
 3 files changed, 52 insertions(+), 13 deletions(-)

diff --git a/shepherd/agent/plugin.py b/shepherd/agent/plugin.py
index f943a80..6451173 100644
--- a/shepherd/agent/plugin.py
+++ b/shepherd/agent/plugin.py
@@ -24,10 +24,38 @@ from .. import base_plugins
 
 log = logging.getLogger(__name__)
 
-# Cache of loaded plugins so far
+# Cache of loaded plugin interfaces so far.
 _loaded_plugins = {}
 
 
+def unload_plugins():
+    """
+    Clear the list of loaded plugins. If the same module is later loaded as a plugin, it will
+    be reloaded.
+    """
+    for plugin_name in _loaded_plugins.copy().keys():
+        unload_plugin(plugin_name)
+
+
+def unload_plugin(plugin_name):
+    """
+    Remove the named plugin from the list of loaded plugins. If the same module is later loaded
+    as a plugin, it will be reloaded. Returns False if the plugin was not already loaded.
+
+    Unloading plugins _should not be relied upon_ to completely reset their state. It is
+    intended primarily for use in testing.
+    Critically, loading a plugin again after unloading it will cause `importlib.reload()` to be
+    called on the primary module or package, _but not its own submodules or other imports_. There
+    is no easy solution to this problem, which is why Shepherd restarts the whole interpreter
+    process to restart.
+    """
+    if plugin_name in _loaded_plugins:
+        del _loaded_plugins[plugin_name]
+        return True
+
+    return False
+
+
 class UnboundMethod():
     """
     Simple wrapper to mark that this is a reference to a method hasn't been bound to an instance
@@ -223,6 +251,9 @@ class InterfaceCall():
         """
         return self.function(**self.kwargs)
 
+    def __repr__(self):
+        return F"{self.plugin_name}.{self.function_name}({self.kwargs})"
+
 
 class InterfaceFunction():
     def __init__(self, func, name=None, remote_command=False):
@@ -582,8 +613,14 @@ def load_plugin(plugin_name, plugin_dir=None):
     # allow them to be listed. Using a try/except block wouldn't be able to tell the difference
     # between a plugin not being found or //it's// imports not loading correctly.
     module = None
+    existing_modules = sys.modules.copy().values()
+
     if plugin_name in discover_base_plugins():
         module = importlib.import_module(base_plugins.__name__+'.'+plugin_name)
+
+        if module in existing_modules:
+            log.info(F"Module for {plugin_name} was aleady imported, reloading")
+            importlib.reload(module)
         log.info(F"Loading base plugin {plugin_name}")
 
     elif plugin_name in discover_custom_plugins(plugin_dir):
@@ -591,6 +628,10 @@ def load_plugin(plugin_name, plugin_dir=None):
         try:
             sys.path = [str(plugin_dir)]
             module = importlib.import_module(plugin_name)
+
+            if module in existing_modules:
+                log.info(F"Module for {plugin_name} was aleady imported, reloading")
+                importlib.reload(module)
         finally:
             sys.path = saved_syspath
         modulepath = getattr(module, "__path__", [module.__file__])[0]
@@ -598,6 +639,10 @@ def load_plugin(plugin_name, plugin_dir=None):
 
     elif plugin_name in discover_installed_plugins():
         module = pkg_resources.iter_entry_points('shepherd.plugins', plugin_name)[0].load()
+
+        if module in existing_modules:
+            log.info(F"Module for {plugin_name} was aleady imported, reloading")
+            importlib.reload(module)
         log.info(F"Loading installed plugin {plugin_name} from {module.__name__}")
 
     if not module:
diff --git a/tests/test_core.py b/tests/test_core.py
index ea3cefd..9b21019 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -5,10 +5,12 @@ import logging
 import pytest
 
 from shepherd.agent import core
+from shepherd.agent import plugin
 
 
 @pytest.fixture
 def local_agent():
+    plugin.unload_plugins()
     return core.Agent(control_enabled=False)
 
 
diff --git a/tests/test_plugins.py b/tests/test_plugins.py
index 11c5e01..168447c 100644
--- a/tests/test_plugins.py
+++ b/tests/test_plugins.py
@@ -7,14 +7,6 @@ import pytest
 from shepherd.agent import plugin
 
 
-def clear_plugin_state(plugin_name):
-    # Make sure it's loading a fresh copy
-    sys.modules.pop(plugin_name, None)
-    plugin._loaded_plugins = {}
-    plugin.plugin_interfaces = {}
-    plugin.plugin_functions = {}
-
-
 @pytest.fixture
 def simple_plugin(request):
     interface = plugin.load_plugin("simpletestplugin", Path(request.fspath.dirname)/'assets')
@@ -34,10 +26,10 @@ def test_simple_interface_function_load(simple_plugin: plugin.PluginInterface):
 
 @pytest.fixture
 def simple_running_plugin(request):
-    clear_plugin_state("simpletestplugin")
+    plugin.unload_plugin("simpletestplugin")
     interface = plugin.load_plugin("simpletestplugin", Path(request.fspath.dirname)/'assets')
     template_config = interface.confspec.get_template()
-    plugin.init_plugins({"simpletestplugin": template_config}, {"ckey": "cval"}, {})
+    plugin.init_plugins({"simpletestplugin": template_config})
     return interface
 
 
@@ -78,10 +70,10 @@ def test_dirty_plugin_load(request):
 
 @pytest.fixture
 def running_class_plugin(request):
-    clear_plugin_state("classtestplugin")
+    plugin.unload_plugin("classtestplugin")
     interface = plugin.load_plugin("classtestplugin", Path(request.fspath.dirname)/'assets')
     template_config = interface.confspec.get_template()
-    plugin.init_plugins({"classtestplugin": template_config}, {"core_conf": "core_conf_vals"}, {})
+    plugin.init_plugins({"classtestplugin": template_config})
     return interface