[tests] PolicyIterationModified unit tests

Two unit tests based on the small example for the discounted and
undiscounted case.

src/tests/test_PolicyIterationModified.py

@@ -5,3 +5,23 @@ Created on Wed Sep 11 13:33:16 2013
 @author: steve
 """
 
+from nose.tools import assert_equal
+import numpy as np
+import scipy.sparse as sp
+
+from mdptoolbox import mdp
+
+from utils import BaseTestMDP, assert_sequence_almost_equal
+
+class TestPolicyIterationModified(BaseTestMDP):
+    def test_small(self):
+        pim = mdp.PolicyIterationModified(self.small_P, self.small_R, 0.9)
+        pim.run()
+        assert_sequence_almost_equal(pim.V,
+                                    (41.8656419239403, 35.4702797722819))
+        assert_equal(pim.policy, (1, 0))
+
+    def test_small_undiscounted(self):
+        pim = mdp.PolicyIterationModified(self.small_P, self.small_R, 1)
+        pim.run()
+        assert_equal(pim.policy, (1, 0))

src/tests/utils.py

@@ -5,6 +5,7 @@ Created on Sat Aug 24 14:44:07 2013
 @author: steve
 """
 
+from nose.tools import assert_true
 import numpy as np
 import scipy as sp
 
@@ -13,6 +14,9 @@ import mdptoolbox.example
 class BaseTestMDP(object):
     small_P, small_R = mdptoolbox.example.small()
 
+def assert_sequence_almost_equal(a, b, spacing=10e-12):
+    assert_true(all(abs(a[k] - b[k]) < spacing for k in range(len(a))))
+
 STATES = 10
 ACTIONS = 3
 SMALLNUM = 10e-12