make it python2 compatible

1 files changed, 9 insertions, 6 deletions

diff --git a/tests/test_pickle.py b/tests/test_pickle.py
index 4bb1d46..6417c86 100644
--- a/tests/test_pickle.py
+++ b/tests/test_pickle.py
@@ -1,5 +1,5 @@
 import unittest
-import array
+import numpy as np
 from pyisda.legs import ContingentLeg, FeeLeg
 from pyisda.curve import SpreadCurve
 from pickle import dumps, loads
@@ -18,7 +18,7 @@ class TestPickle(unittest.TestCase):
 
     def test_yc(self):
         orig_dfs = [self.yc.discount_factor(d) for d in self.yc.dates]
-        pickled_yc = loads(dumps(self.yc))
+        pickled_yc = loads(dumps(self.yc, 2))
         self.assertListAlmostEqual([pickled_yc.discount_factor(d) for d in pickled_yc.dates],
                                    orig_dfs)
 
@@ -30,13 +30,16 @@ class TestPickle(unittest.TestCase):
         cl = ContingentLeg(start_date, end_date, 1)
         fl = FeeLeg(start_date, end_date, True, 1, 1)
         sc = SpreadCurve(self.trade_date, self.yc, start_date,
-                         step_in_date, value_date, [end_date], array.array('d', [75*1e-4]), 0.4)
-        sc_pickled = loads(dumps(sc))
+                         step_in_date, value_date, [end_date], np.array([75*1e-4]), 0.4)
+        sc_pickled = loads(dumps(sc, 2))
         self.assertListAlmostEqual([sc.survival_probability(d) for d in self.yc.dates],
                                    [sc_pickled.survival_probability(d) for d in self.yc.dates])
-        cl_pickled = loads(dumps(cl))
-        fl_pickled = loads(dumps(fl))
+        cl_pickled = loads(dumps(cl, 2))
+        fl_pickled = loads(dumps(fl, 2))
         self.assertEqual(cl_pickled.pv(self.trade_date, step_in_date, value_date, self.yc, sc, 0.4),
                          cl.pv(self.trade_date, step_in_date, value_date, self.yc, sc, 0.4))
         self.assertEqual(fl.pv(self.trade_date, step_in_date, value_date, self.yc, sc, False),
                          fl_pickled.pv(self.trade_date, step_in_date, value_date, self.yc, sc, False))
+
+if __name__=="__main__":
+    unittest.main()