Agent step can be a generator

tests/test_agents.py

@@ -13,14 +13,35 @@ class Dead(agents.FSM):
 
 
 class TestMain(TestCase):
+    def test_die_returns_infinity(self):
+        '''The last step of a dead agent should return time.INFINITY'''
+        d = Dead(unique_id=0, model=environment.Environment())
+        ret = d.step().abs(0)
+        print(ret, "next")
+        assert ret == stime.INFINITY
+
     def test_die_raises_exception(self):
+        '''A dead agent should raise an exception if it is stepped after death'''
         d = Dead(unique_id=0, model=environment.Environment())
         d.step()
         with pytest.raises(agents.DeadAgent):
             d.step()
 
-    def test_die_returns_infinity(self):
-        d = Dead(unique_id=0, model=environment.Environment())
-        ret = d.step().abs(0)
-        print(ret, "next")
-        assert ret == stime.INFINITY
+
+    def test_agent_generator(self):
+        '''
+        The step function of an agent could be a generator. In that case, the state of the
+        agent will be resumed after every call to step.
+        '''
+        class Gen(agents.BaseAgent):
+            def step(self):
+                a = 0
+                for i in range(5):
+                    yield a
+                    a += 1
+        e = environment.Environment()
+        g = Gen(model=e, unique_id=e.next_id())
+
+        for i in range(5):
+            t = g.step()
+            assert t == i