Tests for attaching hooks (generic)

dwhswenson · Jun 29, 2021 · Jan 29, 2018 · Jan 29, 2018 · Jan 29, 2018 · Jan 30, 2018
commit 64639b71c7d255a4479b3217c006881552944891
diff --git a/openpathsampling/pathsimulators/path_simulator.py b/openpathsampling/pathsimulators/path_simulator.py
@@ -68,6 +68,8 @@ class PathSimulator(with_metaclass(abc.ABCMeta, StorableNamedObject)):
     calc_name = "PathSimulator"
     _excluded_attr = ['sample_set', 'step', 'save_frequency',
                       'output_stream']
+    hook_names = ['before_simulation', 'before_step', 'after_step',
+                  'after_simulation']
 
     def __init__(self, storage):
         super(PathSimulator, self).__init__()
@@ -82,8 +84,7 @@ def __init__(self, storage):
         self.sample_set = None
         self.output_stream = sys.stdout  # user can change to file handler
         self.allow_refresh = True
-        self.hooks = {k: [] for k in ['before_simulation', 'before_step',
-                                      'after_step', 'after_simulation']}
+        self.hooks = self.empty_hooks()
 
     def sync_storage(self):
         """
@@ -92,14 +93,39 @@ def sync_storage(self):
         if self.storage is not None:
             self.storage.sync_all()
 
+    def empty_hooks(self):
+        """Return a hook dictionary with no hooks."""
+        return {k: [] for k in self.hook_names}
+
     def attach_hook(self, hook, hook_for=None):
+        """Attach a hook class or method to this simulation.
+
+        Parameters
+        ----------
+        hook : :class:`.PathSimulatorHook` or method
+            Hook to add
+        hook_for : str or None
+            If None (default) then the ``hook`` must be a class with methods
+            named to match the hook names in this simulator. If ``hook`` is
+            a method, then ``hook_for`` must be the name of the hook it
+            represents
+        """
+        def add_hook_method(hook_method, hook_name):
+            try:
+                self.hooks[hook_name].append(hook_method)
+            except KeyError:
+                raise TypeError("No hook '" + hook_name + "' in " +
+                                str(self.__class__.__name__))
+
         if hook_for is None:
             for hook_name in hook.implemented_for:
-                self[hook_name].append(getattr(hook, hook_name))
+                hook_method = getattr(hook, hook_name)
+                add_hook_method(hook_method, hook_name)
         else:
-            self[hook_for].append(hook)
+            add_hook_method(hook, hook_for)
 
     def run_hooks(self, hook_name, **kwargs):
+        """Run the hooks for the given ``hook_name``"""
         for hook in self.hooks[hook_name]:
             hook(**kwargs)
 

diff --git a/openpathsampling/pathsimulators/shoot_snapshots.py b/openpathsampling/pathsimulators/shoot_snapshots.py
@@ -17,6 +17,17 @@ class ShootFromSnapshotsSimulation(PathSimulator):
     user can create a simulation of this sort on-the-fly for some weird
     ensembles.
 
+    .. note::
+
+        **Hooks**: This can interact with the :class:`.PathSimulatorHook`
+        framework. In these hooks, ``step_info`` is the 4-tuple
+        ``(snap, n_snaps, step, n_steps)`` where ``snap`` is the snapshot
+        number, ``step`` is the step number within that snapshot, and
+        ``n_snaps`` and ``n_steps`` are the total number of each,
+        respectively. The ``state`` will provide the initial snapshot we are
+        shooting from, and ``results`` is the :class:`.MCStep` that comes
+        out of each step.
+
     Parameters
     ----------
     storage : :class:`.Storage`

diff --git a/openpathsampling/tests/pathsimulators/test_hooks.py b/openpathsampling/tests/pathsimulators/test_hooks.py
@@ -18,32 +18,147 @@
 logging.getLogger('openpathsampling.storage').setLevel(logging.CRITICAL)
 logging.getLogger('openpathsampling.netcdfplus').setLevel(logging.CRITICAL)
 
+class TrivialPathSimulator(paths.PathSimulator):
+    """
+    use this to check that the hooks get attached and run correctly
+    """
+    def run(self, n_steps):
+        self.run_hooks('before_simulation', sim=self)
+        state = ""
+        for step in range(n_steps):
+            self.run_hooks('before_step', sim=self, step_number=step,
+                           step_info=(step, n_steps), state=state)
+            result = str(step)
+            state += result
+            self.run_hooks('after_step', sim=self, step_number=step,
+                           step_info=(step, n_steps), state=state,
+                           results=result)
+        self.run_hooks('after_simulation', sim=self)
+
+class NoStepHookPathSimulator(paths.PathSimulator):
+    """use this to check that only correct hooks get assigned"""
+    hook_names = ['before_simulation', 'extra_hook']
+    def run(self, n_steps):
+        self.run_hooks('before_simulation', sim=self)
+        self.run_hooks('extra_hook', sim=self, extra="foo")
+
+
+class StupidHook(PathSimulatorHook):
+    def __init__(self):
+        self.began_simulation = False
+        self.began_steps = 0
+        self.finished_steps = 0
+        self.finished_simulation = False
+
+    def before_simulation(self, sim):
+        self.began_simulation = True
+
+    def before_step(self, sim, step_number, step_info, state):
+        self.began_steps +=1
+
+    def after_step(self, sim, step_number, step_info, state, results):
+        self.finished_steps += 1
+
+    def after_simulation(self, sim):
+        self.finished_simulation = True
+
+class StupiderHook(PathSimulatorHook):
+    def __init__(self):
+        self.foo = None
+
+    def hook_method(self, sim, extra=None):
+        self.foo = extra
+
 class TestPathSimulatorHook(object):
-    def setup(self):
-        pass
+    # includes tests for attach_hook and run_hook
+    def test_attach_hooks_class(self):
+        hook = StupidHook()
+        sim = TrivialPathSimulator(storage=None)
+        sim.attach_hook(hook)
+        for k in sim.hook_names:
+            assert_equal(len(sim.hooks[k]), 1)
 
-    def test_attach_hooks(self):
-        pass
+    def test_attach_hooks_method(self):
+        hook_instance = StupiderHook()
+        sim = NoStepHookPathSimulator(storage=None)
+        sim.attach_hook(hook_instance.hook_method, hook_for='extra_hook')
+        assert_not_equal(hook_instance.foo, "foo")
+        sim.run(1)
+        assert_equal(hook_instance.foo, "foo")
+
+    def test_trivial_simulation_hooks(self):
+        sim = TrivialPathSimulator(storage=None)
+        stupid = StupidHook()
+        stupider = StupiderHook()
+        sim.attach_hook(stupid)
+        extra = lambda sim: stupider.hook_method(sim, "foo")
+        sim.attach_hook(extra, hook_for='after_simulation')
+        # before running, check that state is as expected
+        assert_equal(stupid.began_simulation, False)
+        assert_equal(stupid.finished_simulation, False)
+        assert_equal(stupid.began_steps, 0)
+        assert_equal(stupid.finished_steps, 0)
+        assert_not_equal(stupider.foo, "foo")
+        # run
+        sim.run(3)
+        # check results
+        assert_equal(stupid.began_simulation, True)
+        assert_equal(stupid.finished_simulation, True)
+        assert_equal(stupid.began_steps, 3)
+        assert_equal(stupid.finished_steps, 3)
+        assert_equal(stupider.foo, "foo")
+
+    @raises(TypeError)
+    def test_bad_extra_hook_method(self):
+        sim = NoStepHookPathSimulator(storage=None)
+        stupid = StupidHook()
+        sim.attach_hook(stupid)
+
+    @raises(TypeError)
+    def test_bad_hook_name(self):
+        sim = NoStepHookPathSimulator(storage=None)
+        stupider = StupiderHook()
+        sim.attach_hook(stupider.hook_method, "blah")
+
+    def test_no_step_hook_simulation_hooks(self):
+        sim = NoStepHookPathSimulator(storage=None)
+        stupid = StupidHook()
+        stupider = StupiderHook()
+        sim.attach_hook(stupid.before_simulation,
+                        hook_for='before_simulation')
+        sim.attach_hook(stupider.hook_method, hook_for='extra_hook')
+        assert_equal(len(sim.hooks), len(sim.hook_names))
+        for hook in sim.hooks:
+            assert_equal(len(sim.hooks[hook]), 1)
+
+        assert_equal(stupid.began_simulation, False)
+        assert_not_equal(stupider.foo, "foo")
+        sim.run(1)
+        assert_equal(stupid.began_simulation, True)
+        assert_equal(stupider.foo, "foo")
 
 class TestStorageHook(object):
+    def _delete_files(self):
+        pass
+
     def setup(self):
         pass
 
     def test_before_simulation(self):
-        pass
+        raise SkipTest
 
     def test_after_step(self):
-        pass
+        raise SkipTest
 
     def test_after_simulation(self):
-        pass
+        raise SkipTest
 
 class TestShootFromSnapshotsOutputHook(object):
     def setup(self):
         pass
 
     def test_before_simulation(self):
-        pass
+        raise SkipTest
 
     def test_before_step(self):
-        pass
+        raise SkipTest