From 2a7b950e57f241e2c14ee48b3932fd7603eb1342 Mon Sep 17 00:00:00 2001
From: novirium <t.wilson.au@gmail.com>
Date: Mon, 1 Feb 2021 16:46:10 +0800
Subject: [PATCH] Start restructuring Core to start more cleanly

---
 setup.py                                     |   5 +-
 shepherd/agent/cli.py                        | 121 +++++++++++++------
 shepherd/agent/control.py                    |   2 +-
 shepherd/agent/core.py                       |  57 +++++----
 shepherd/agent/plugin.py                     |   4 +-
 shepherd/agent/tasks.py                      |  16 ++-
 tests/assets/shepherd-test.toml              |   4 +
 tests/assets/shepherd-test2.toml             |   7 ++
 tests/assets/testpluginpackage/__init__.py   |   8 ++
 tests/assets/testpluginpackage/testplugin.py |  32 +++++
 tests/test_cli.py                            |  59 +++++++++
 tests/test_control.py                        |   2 +-
 tests/test_core.py                           |   2 +-
 13 files changed, 244 insertions(+), 75 deletions(-)
 create mode 100644 tests/assets/shepherd-test.toml
 create mode 100644 tests/assets/shepherd-test2.toml
 create mode 100644 tests/assets/testpluginpackage/__init__.py
 create mode 100644 tests/assets/testpluginpackage/testplugin.py
 create mode 100644 tests/test_cli.py

diff --git a/setup.py b/setup.py
index 5121341..056aa6d 100644
--- a/setup.py
+++ b/setup.py
@@ -14,7 +14,8 @@ setup(
         'apscheduler',
         'paramiko',
         'python-dateutil',
-        'click'
+        'click',
+        'chromalog'
     ],
     extras_require={
         'dev': [
@@ -35,7 +36,7 @@ setup(
     },
 
     entry_points={
-        'console_scripts': ['shepherd=shepherd.agent.core:cli'],
+        'console_scripts': ['shepherd=shepherd.agent.cli:cli'],
     },
     license='GPLv3+',
     description='Herd your mob of physically remote nodes',
diff --git a/shepherd/agent/cli.py b/shepherd/agent/cli.py
index bac403e..11d5191 100644
--- a/shepherd/agent/cli.py
+++ b/shepherd/agent/cli.py
@@ -3,6 +3,7 @@
 import logging
 import os
 import sys
+from pathlib import Path
 import glob
 from types import SimpleNamespace
 from datetime import datetime
@@ -11,12 +12,13 @@ import pkg_resources
 
 import chromalog
 import click
+import toml
 
-from . import core
+from . import core, plugin
 
-chromalog.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))
+# chromalog.basicConfig(level=os.environ.get("LOGLEVEL", "INFO"))
 
-log = logging.getLogger("shepherd.agent")
+log = logging.getLogger("shepherd.cli")
 
 
 def echo_heading(title, on_nl=True):
@@ -46,22 +48,25 @@ def echo_section(title, input_text=None, on_nl=True):
               " Shepherd Control remote features")
 @click.option('-d', '--default-config-only', 'only_default_layer', is_flag=True,
               help="Ignore the custom config layer (still uses the Control config above that)")
-@click.option('-n', '--new', 'new_run', is_flag=True,
+@click.option('-n', '--new-device-mode', 'new_device_mode', is_flag=True,
               help="Clear existing device identity and cached Shepherd Control config layer."
               " Also triggered by the presence of a shepherd.new file in the"
               " same directory as the custom config layer file.")
 @click.pass_context
-def cli(ctx, default_config_path, local_operation, only_default_layer, new_run):
+def cli(ctx, default_config_path, local_operation, only_default_layer, new_device_mode):
     """
     Core service. If default config file is not provided with '-c' option, the first filename
     in the current working directory beginning  with "shepherd" and
     ending with ".toml" will be used.
     """
-    ctx.ensure_object(SimpleNamespace)
 
     version_text = pkg_resources.get_distribution("shepherd")
     log.info(F"Shepherd Agent [{version_text}]")
 
+    agent = core.Agent()
+    ctx.ensure_object(SimpleNamespace)
+    ctx.obj.agent = agent
+
     # Drop down to subcommand if it doesn't need default config file processing
     if ctx.invoked_subcommand == "template":
         return
@@ -84,34 +89,25 @@ def cli(ctx, default_config_path, local_operation, only_default_layer, new_run):
             sys.exit(1)
 
     # Establish what config layers we're going to try and use
-    layers_disabled = []
+    control_enabled = True
+    use_custom_config = True
+
     if local_operation or (ctx.invoked_subcommand == "test"):
-        layers_disabled.append("control")
+        control_enabled = False
         log.info("Running in local only mode")
 
     if only_default_layer:
-        layers_disabled.append("custom")
+        use_custom_config = False
 
-    agent = core.Agent(control_enabled)
-    agent.load(default_config_path, use_custom_config, new_device_mode)
+    agent.load(default_config_path, use_custom_config, control_enabled, new_device_mode)
 
     # Drop down to subcommands that needed a config compiled
     if ctx.invoked_subcommand == "test":
-        ctx.obj.agent = agent
         return
 
-    agent.start()
-
     print(str(datetime.now()))
 
-    if ctx.invoked_subcommand is not None:
-        return
-
-    print('Press Ctrl+{0} to exit'.format('Break' if os.name == 'nt' else 'C'))
-    try:
-        scheduler.start()
-    except (KeyboardInterrupt, SystemExit):
-        pass
+    agent.start()
 
 
 @cli.command()
@@ -119,30 +115,37 @@ def cli(ctx, default_config_path, local_operation, only_default_layer, new_run):
 @click.argument('interface_function', required=False)
 @click.pass_context
 def test(ctx, plugin_name, interface_function):
-    echo_heading("Shepherd Test")
+    agent = ctx.obj.agent
+    plugin_configs = agent.applied_config.copy()
+    del plugin_configs['shepherd']
+
+    echo_heading("Shepherd - Test")
 
     if not plugin_name:
+        log.info("Test initialisation of all plugins in config...")
 
         echo_section("Plugins loaded:")
-        if len(ctx.obj.plugin_configs) == 0:
+        if len(plugin_configs) == 0:
             click.echo("---none---")
-        for plugin_name, config in ctx.obj.plugin_configs.items():
+        for plugin_name, config in plugin_configs.items():
             click.secho(F" {plugin_name}", fg='green')
 
         echo_section("Core config:")
-        pprint(ctx.obj.core_config)
+        print(toml.dumps(agent.core_config))
+        # pprint(agent.core_config)
 
         echo_section("Plugin configs:")
-        if len(ctx.obj.plugin_configs) == 0:
+        if len(plugin_configs) == 0:
             click.echo("---none---")
-        for name, config in ctx.obj.plugin_configs.items():
+        for name, config in plugin_configs.items():
             click.secho(F" {plugin_name}", fg='green')
-            pprint(config)
+            print(toml.dumps(config))
+            # pprint(config)
 
         click.echo("")
 
         log.info("Initialising plugins...")
-        plugin.init_plugins(ctx.obj.plugin_configs, ctx.obj.core_config)
+        plugin.init_plugins(agent.applied_config)
         log.info("Plugin initialisation done")
 
         return
@@ -150,34 +153,61 @@ def test(ctx, plugin_name, interface_function):
     echo_section("Target plugin:", input_text=plugin_name, on_nl=False)
     # TODO find plugin dependancies
 
-    if plugin_name not in ctx.obj.plugin_configs:
+    if plugin_name not in plugin_configs:
         log.error(F"Supplied plugin name '{plugin_name}' is not loaded"
                   " (not present in config)")
         sys.exit(1)
 
     echo_section(F"Config [{plugin_name}]:")
-    pprint(ctx.obj.plugin_configs[plugin_name])
+    print(toml.dumps(plugin_configs[plugin_name]))
+    # pprint(plugin_configs[plugin_name])
 
     interface = plugin.load_plugin(plugin_name)
 
     if not interface_function:
-        echo_section(F"Interface functions [{plugin_name}]:")
+        echo_section(F"Interface functions [{plugin_name}]:", on_nl=False)
 
         for name, func in interface._functions.items():
             click.echo(F" {name}")
         return
 
-    echo_section("Target interface function:", input_text=interface_function)
+    echo_section("Target interface function:", input_text=interface_function,  on_nl=False)
     if interface_function not in interface._functions:
         log.error(F"Supplied interface function name '{interface_function}' is not present in"
                   F" plugin {plugin_name}")
         sys.exit(1)
 
     log.info("Initialising plugins...")
-    plugin.init_plugins({plugin_name: ctx.obj.plugin_configs[plugin_name]}, ctx.obj.core_config)
+    plugin.init_plugins({plugin_name: plugin_configs[plugin_name]})
     log.info("Plugin initialisation done")
 
-    interface._functions[interface_function]()
+    # TODO look for a spec on the interface function, and parse cmdline values if it's there
+
+    print(interface._functions[interface_function]())
+
+
+class BlankEncoder(toml.TomlEncoder):
+    """
+    A TOML encoder that emit empty keys (values of None). This isn't valid TOML,
+    but is useful for generating templates.
+    """
+    class BlankValue:
+        pass
+
+    def __init__(self, _dict=dict, preserve=False):
+        super().__init__(_dict, preserve)
+        self.dump_funcs[self.BlankValue] = lambda v: ''
+
+    def dump_sections(self, o, sup):
+        for section in o:
+            if o[section] is None:
+                o[section] = self.BlankValue()
+        return super().dump_sections(o, sup)
+
+    def dump_value(self, v):
+        if v is None:
+            v = self.BlankValue()
+        return super().dump_value(v)
 
 
 @cli.command()
@@ -199,13 +229,17 @@ def template(ctx, plugin_name, include_all, config_path, plugin_dir):
     a new file (if it doesn't yet exist).
     """
 
+    agent = ctx.obj.agent
+
+    echo_heading("Shepherd - Template")
+
     if not plugin_dir:
         plugin_dir = Path.cwd()
 
-    confspec = ConfigSpecification()
+    confspec = None
     if (not plugin_name) or (plugin_name == "shepherd"):
         plugin_name = "shepherd"
-        confspec = core_confspec()
+        confspec = agent.core_interface.confspec
     else:
         try:
             plugin_interface = plugin.load_plugin(plugin_name, plugin_dir)
@@ -215,9 +249,16 @@ def template(ctx, plugin_name, include_all, config_path, plugin_dir):
         confspec = plugin_interface.confspec
 
     template_dict = confspec.get_template(include_all)
-    template_toml = toml.dumps({plugin_name: template_dict})
+    template_toml = toml.dumps({plugin_name: template_dict}, encoder=BlankEncoder())
 
-    log.info(F"Config template for [{plugin_name}]: \n\n"+template_toml)
+    if include_all:
+        log.info("Including all optional fields")
+    else:
+        log.info("Including required fields only")
+
+    echo_section("Config template for", input_text=F"[{plugin_name}]")
+    click.echo("")
+    click.echo(template_toml)
 
     if not config_path:
         # reuse parent "-c" for convenience
diff --git a/shepherd/agent/control.py b/shepherd/agent/control.py
index 0b9733e..e55c38b 100644
--- a/shepherd/agent/control.py
+++ b/shepherd/agent/control.py
@@ -30,7 +30,7 @@ def _update_required_callback():
         _control_update_required.notify()
 
 
-def register(core_interface):
+def register_on(core_interface):
     """
     Register the control confspec on the core interface.
     """
diff --git a/shepherd/agent/core.py b/shepherd/agent/core.py
index a04432b..698919b 100644
--- a/shepherd/agent/core.py
+++ b/shepherd/agent/core.py
@@ -24,7 +24,7 @@ class Agent():
     Holds the main state required to run Shepherd Agent
     """
 
-    def __init__(self, control_enabled=True):
+    def __init__(self):
         # The config defined by the device (everything before the Control layer)
         self.local_config = None
         # The config actually being used
@@ -33,11 +33,23 @@ class Agent():
         self.core_config = None
 
         self.interface_functions = None
+        self.control_enabled = None
+        self.plugin_interfaces = None
 
-        self.control_enabled = control_enabled
+        self.restart_args = None
 
+        # Setup core interface
         self.core_interface = plugin.PluginInterface()
-        self.plugin_interfaces = None
+        self.core_interface.register_confspec(core_confspec())
+        self.core_interface.register_function(self.root_dir)
+        self.core_interface.register_function(self.device_name)
+        # Allows plugins to add delay for system time to stabilise
+        self.core_interface.register_hook("wait_for_stable_time")
+
+        # Allow other modules to add to the core interface (confspec, hooks, interface functions)
+        # Having modules modify a confspec after it's registered here is a bit of a hack.
+        tasks.register_on(self.core_interface)
+        control.register_on(self.core_interface)
 
     def root_dir(self):
         return self.core_config["root_dir"]
@@ -45,7 +57,8 @@ class Agent():
     def device_name(self):
         return self.core_config["name"]
 
-    def load(self, default_config_path, use_custom_config=True,  new_device_mode=False):
+    def load(self, default_config_path, use_custom_config=True, control_enabled=False,
+             new_device_mode=False):
         """
         Load in the Shepherd Agent config and associated plugins.
         Args:
@@ -57,18 +70,10 @@ class Agent():
                 of ID, as if it were being run on a fresh system.
         """
 
-        # Setup core interface
-        self.core_interface.register_confspec(core_confspec())
-        self.core_interface.register_function(self.root_dir)
-        self.core_interface.register_function(self.device_name)
-        # Allows plugins to add delay for system time to stabilise
-        self.core_interface.register_hook("wait_for_stable_time")
-
-        # Allow other modules to add to the core interface (confspec, hooks, interface functions)
-        # Having modules modify a confspec after it's registered here is a bit of a hack.
-        tasks.register(self.core_interface)
-        control.register(self.core_interface)
+        self.restart_args = [default_config_path,
+                             use_custom_config, control_enabled, new_device_mode]
 
+        self.control_enabled = control_enabled
         # Because the plugin module caches interfaces, this will then get used when loading
         # config layers and validating them
         plugin.load_plugin_interface("shepherd", self.core_interface)
@@ -105,6 +110,12 @@ class Agent():
         self.applied_config = confman.get_config_bundles()
         self.core_config = confman.get_config_bundle('shepherd')
 
+        loaded_plugin_names = list(self.applied_config.keys())
+        loaded_plugin_names.remove('shepherd')
+        if len(loaded_plugin_names) == 0:
+            loaded_plugin_names.append("--none--")
+        log.info(F"Loaded plugins: {', '.join(loaded_plugin_names)}")
+
         log.debug("Compiled config: %s", confman.root_config)
         if self.core_config["compiled_config_path"]:
             message = F"Compiled Shepherd config at {datetime.now()}"
@@ -115,10 +126,11 @@ class Agent():
         pass
 
     def start(self):
-
-        # After this point, plugins may already have their own threads running if they create
-        # them during init
+        # We don't worry about the plugin dir here, or 'shepherd' being included, as they should
+        # already all be loaded and cached.
         self.plugin_interfaces = plugin.init_plugins(self.applied_config)
+        # After this point, plugins may already have their own threads running if they created
+        # them during init
         self.interface_functions = self.core_interface.plugins
 
         cmd_runner = control.CommandRunner(self.interface_functions)
@@ -135,9 +147,6 @@ class Agent():
 
         # Need somewhere to eventually pass in the hooks Tasks will need for the lowpower stuff,
         # probably just another init_plugins arg.
-        # Eventually when the dust settles we might revisit converting the core "shepherd"
-        # namespace stuff into it's own plugin interface, as it's using a lot of the same
-        # mechanisms, but we're having to pass it all around individually.
 
         # TODO Collect plugin tasks
 
@@ -247,13 +256,13 @@ def compile_remote_config(confman):
         control_config = control.get_cached_config(core_conf["root_dir"])
         try:
             load_config_layer_and_plugins(confman, control_config)
-            log.info(F"Loaded cached Shepherd Control config layer")
+            log.info("Loaded cached Shepherd Control config layer")
         except Exception as e:
             if isinstance(e, InvalidConfigError):
-                log.error(F"Failed to load cached Shepherd Control config layer."
+                log.error("Failed to load cached Shepherd Control config layer."
                           F" {e.args[0]}")
             else:
-                log.error(F"Failed to load cached Shepherd Control config layer.",
+                log.error("Failed to load cached Shepherd Control config layer.",
                           exc_info=True)
             log.warning("Falling back to local config.")
             confman.fallback()
diff --git a/shepherd/agent/plugin.py b/shepherd/agent/plugin.py
index 6451173..40a20d2 100644
--- a/shepherd/agent/plugin.py
+++ b/shepherd/agent/plugin.py
@@ -284,7 +284,7 @@ class InterfaceFunction():
 
                 if not isinstance(arg_spec, _ValueSpecification):
                     raise ValueError("Function annotations for a Shepherd Interface function"
-                                     "must be a type of ConfigSpecification, or on the the valid"
+                                     "must be a type of ConfigSpecification, or one of the valid"
                                      "type shortcuts")
 
                 self.spec.add_spec(param.name, arg_spec)
@@ -298,7 +298,7 @@ class InterfaceFunction():
 
                 if not isinstance(ret_spec, _ValueSpecification):
                     raise ValueError("Function annotations for a Shepherd Interface function"
-                                     "must be a type of ConfigSpecification, or on the the valid"
+                                     "must be a type of ConfigSpecification, or one of the valid"
                                      "type shortcuts")
 
                 self.spec.add_spec("return", arg_spec)
diff --git a/shepherd/agent/tasks.py b/shepherd/agent/tasks.py
index 7df0310..121997c 100644
--- a/shepherd/agent/tasks.py
+++ b/shepherd/agent/tasks.py
@@ -23,7 +23,7 @@ from .util import HoldLock
 log = logging.getLogger("shepherd.agent.tasks")
 
 
-def register(core_interface):
+def register_on(core_interface):
     """
     Register the session confspec and hooks on the core interface passed in - `start_tasks` later
     assumes that these hooks are present.
@@ -36,7 +36,7 @@ def register(core_interface):
     confspec.add_spec("min_suspend_time", IntSpec(helptext="Smallest wait period before the next"
                                                   " scheduled task that the agent will decide to"
                                                   " suspend, in seconds"), default=300)
-    
+
     core_interface.confspec.add_spec("session", confspec, optional=True, default={})
 
     # `resume_time` is a DateTime indicating when the session should resume - it already has the
@@ -256,6 +256,7 @@ MIN_DELAY = 0.01  # Minimum time (in seconds) the task loop will sleep for.
 def _tasks_update_loop(config, suspend_hook, session):
 
     sched_tasks = []
+    # When resuming, schedule tasks from the desired resume time, even if it's in the past
     base_time = session.resume_time
     now = datetime.now(tz.tzutc())
     # If it's a new session, only schedule tasks from now.
@@ -282,8 +283,13 @@ def _tasks_update_loop(config, suspend_hook, session):
         sched_tasks.append(ScheduledTask(scheduled_time, task))
 
     suspend_available = False
-    if config['enable_suspend'] and suspend_hook.attachments:
-        suspend_available = True
+    if config['enable_suspend']:
+        if suspend_hook.attachments:
+            suspend_available = True
+            log.info("Session suspension enabled.")
+        else:
+            log.warning("'enable_suspend' set to true, but no suspend hooks are attached. Add"
+                        " a plugin that provides a suspend hook.")
 
     # Let our `start_tasks` call continue
     _update_thread_init_done.set()
@@ -296,6 +302,8 @@ def _tasks_update_loop(config, suspend_hook, session):
         if sched_tasks[0].scheduled_for <= now:
             # Scheduled time has passed, run the task
             log.info(F"Running task {sched_tasks[0].task.interface_call}...")
+
+            # Should we be catching exceptions for this?
             sched_tasks[0].task.interface_call.call()
 
             # Reschedule and sort
diff --git a/tests/assets/shepherd-test.toml b/tests/assets/shepherd-test.toml
new file mode 100644
index 0000000..6662567
--- /dev/null
+++ b/tests/assets/shepherd-test.toml
@@ -0,0 +1,4 @@
+[shepherd]
+name = "shepherd-test"
+root_dir ="./"
+compiled_config_path = ""
diff --git a/tests/assets/shepherd-test2.toml b/tests/assets/shepherd-test2.toml
new file mode 100644
index 0000000..1a9f566
--- /dev/null
+++ b/tests/assets/shepherd-test2.toml
@@ -0,0 +1,7 @@
+[shepherd]
+name = "shepherd-test"
+root_dir ="./"
+compiled_config_path = ""
+plugin_dir = "./"
+[classtestplugin]
+spec1 = "a"
diff --git a/tests/assets/testpluginpackage/__init__.py b/tests/assets/testpluginpackage/__init__.py
new file mode 100644
index 0000000..e05f84b
--- /dev/null
+++ b/tests/assets/testpluginpackage/__init__.py
@@ -0,0 +1,8 @@
+from configspec import *
+from shepherd import PluginInterface, plugin, plugin_function, plugin_hook, plugin_attachment
+
+interface = PluginInterface()
+
+
+confspec = ConfigSpecification()
+confspec.add_spec("spec2", StringSpec(helptext="helping!"))
\ No newline at end of file
diff --git a/tests/assets/testpluginpackage/testplugin.py b/tests/assets/testpluginpackage/testplugin.py
new file mode 100644
index 0000000..7265a58
--- /dev/null
+++ b/tests/assets/testpluginpackage/testplugin.py
@@ -0,0 +1,32 @@
+
+
+
+@plugin
+class SystemPlugin():
+    def __init__(self, pluginInterface, config):
+        super().__init__(pluginInterface, config)
+        self.config = config
+        self.interface = pluginInterface
+        self.plugins = pluginInterface.other_plugins
+        self.hooks = pluginInterface.hooks
+
+        self.interface.register_function(self.echo)
+        self.interface.register_function(self.exec)
+
+    @plugin_function()
+    def echo(self, string: str):
+        pass
+
+    def exec(self):
+        pass
+
+    @plugin_hook
+    def callback(self):
+        pass
+
+    @plugin_attachment("pluginname.hookname")
+    def caller(self):
+        pass
+
+
+# interface.register_plugin(SystemPlugin)
diff --git a/tests/test_cli.py b/tests/test_cli.py
new file mode 100644
index 0000000..99ee537
--- /dev/null
+++ b/tests/test_cli.py
@@ -0,0 +1,59 @@
+# pylint: disable=redefined-outer-name
+from pathlib import Path
+import logging
+
+from click.testing import CliRunner
+import pytest
+
+from shepherd.agent.cli import cli
+
+
+def test_shepherd_template():
+    # Note that the CliRunner doesn't catch log output
+    runner = CliRunner()
+    result = runner.invoke(cli, ['template'])
+    assert """
+.: Shepherd - Template :.
+
+:: Config template for [shepherd]
+
+[shepherd]
+name =""" in result.output
+
+
+def test_shepherd_optional_template():
+    runner = CliRunner()
+    result = runner.invoke(cli, ['template', '-a'])
+    assert """
+.: Shepherd - Template :.
+
+:: Config template for [shepherd]
+
+[shepherd]
+name = 
+root_dir = "./"
+custom_config_path = 
+compiled_config_path = "compiled-config.toml"
+plugin_dir = "./shepherd-plugins"
+
+[shepherd.session]
+resume_delay = 180
+enable_suspend = true
+min_suspend_time = 300
+
+[shepherd.control]
+server = 
+intro_key =""" in result.output
+
+
+def test_plugin_template(request):
+    plugindir = Path(request.fspath.dirname)/'assets'
+    runner = CliRunner()
+    result = runner.invoke(cli, ['template', '-d', str(plugindir), 'simpletestplugin'])
+    assert """
+.: Shepherd - Template :.
+
+:: Config template for [simpletestplugin]
+
+[simpletestplugin]
+spec1 =""" in result.output
diff --git a/tests/test_control.py b/tests/test_control.py
index 2aa46ac..a7c4459 100644
--- a/tests/test_control.py
+++ b/tests/test_control.py
@@ -41,7 +41,7 @@ def control_config():
 def registered_interface():
     interface = plugin.PluginInterface()
     interface.register_confspec(ConfigSpecification())
-    control.register(interface)
+    control.register_on(interface)
     return interface
 
 
diff --git a/tests/test_core.py b/tests/test_core.py
index 9b21019..df38196 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -11,7 +11,7 @@ from shepherd.agent import plugin
 @pytest.fixture
 def local_agent():
     plugin.unload_plugins()
-    return core.Agent(control_enabled=False)
+    return core.Agent()
 
 
 @pytest.fixture