add tests

dwhswenson · Aug 23, 2016 · Aug 13, 2016 · Aug 13, 2016 · Aug 13, 2016 · Aug 14, 2016
commit 5e0f6e6d7287df67d58a3e8cffef70b3d70cf959
diff --git a/openpathsampling/engines/dynamics_engine.py b/openpathsampling/engines/dynamics_engine.py
@@ -430,6 +430,7 @@ def iter_generate(self, initial, running=None, direction=+1,
         valid = False
         attempt_nan = 0
         attempt_error = 0
+        attempt_max_length = 0
         trajectory = initial
 
         final_error = None
@@ -528,13 +529,23 @@ def iter_generate(self, initial, running=None, direction=+1,
                             'Hit maximal length of %d frames.' %
                             self.options['n_frames_max']
                         )
+                        break
                     elif on == 'ignore':
                         # keep on going
                         pass
                     elif on == 'stop':
                         logger.info('Trajectory hit max length. Stopping.')
                         # fail gracefully
                         stop = True
+                    elif on == 'retry':
+                        attempt_max_length += 1
+                        if attempt_max_length > self.retries_when_max_length:
+                            if self.on_nan == 'fail':
+                                final_error = RuntimeError(
+                                    'Failed to generate trajectory without '
+                                    'hitting max length after %d attempts' %
+                                    attempt_max_length)
+                                break
 
             if has_nan:
                 attempt_nan += 1

diff --git a/openpathsampling/tests/test_helpers.py b/openpathsampling/tests/test_helpers.py
@@ -85,10 +85,8 @@ class CalvinistDynamics(DynamicsEngine):
     def __init__(self, predestination):
         topology = Topology(n_atoms=1, n_spatial=1)
         engine = peng.tools.TopologyEngine(topology)
-        template = toys.Snapshot(engine=engine)
 
-        super(CalvinistDynamics, self).__init__(options={'n_frames_max' : 12},
-                                                template=template)
+        super(CalvinistDynamics, self).__init__(options={'n_frames_max': 12})
         self.predestination = make_1d_traj(coordinates=predestination,
                                            velocities=[1.0]*len(predestination),
                                            engine=engine

diff --git a/openpathsampling/tests/testopenmmengine.py b/openpathsampling/tests/testopenmmengine.py
@@ -147,7 +147,16 @@ def test_fail_length(self):
         self.engine.options['n_max_length'] = 2
         _ = self.engine.generate(self.engine.current_snapshot, [true_func])
 
-    # @raises_with_message_like(RuntimeError,
+    @raises_with_message_like(RuntimeError,
+                              'Failed to generate trajectory without hitting '
+                              'max length')
+    def test_retry_length(self):
+        self.engine.on_max_length = 'retry'
+        self.engine.options['n_max_length'] = 2
+        self.engine.retries_when_max_length = 2
+        _ = self.engine.generate(self.engine.current_snapshot, [true_func])
+
+            # @raises_with_message_like(RuntimeError,
     #                           "Failed to generate trajectory without error")
     # def test_fail_length(self):
     #     self.engine.on_nan = 'fail'