more pickle tests

1 files changed, 46 insertions, 8 deletions

diff --git a/tests/test_pickle.py b/tests/test_pickle.py
index cdb698d..476e834 100644
--- a/tests/test_pickle.py
+++ b/tests/test_pickle.py
@@ -1,15 +1,53 @@
 import unittest
-from pyisda.legs import ContingentLeg
+import array
+from pyisda.legs import ContingentLeg, FeeLeg
+from pyisda.curve import SpreadCurve
 from pickle import dumps, loads
 import datetime
+from quantlib.settings import Settings
+from quantlib.time.api import Date
+import sys
+sys.path.append('/home/share/CorpCDOs/code/python/')
+from yieldcurve import YC, get_dates, ql_to_jp
 
 class TestPickle(unittest.TestCase):
 
+    settings = Settings()
+    trade_date = datetime.date(2016, 10, 6)
+    settings.evaluation_date =  Date(6, 10, 2016)
+    yc = YC()
+    jp_yc = ql_to_jp(yc)
+
+    def assertListAlmostEqual(self, l1, l2):
+        if len(l1) != len(l2):
+            raise AssertionError
+        for a, b in zip(l1, l2):
+            self.assertAlmostEqual(a, b)
+
+    def test_yc(self):
+        jp_yc = ql_to_jp(self.yc)
+        dates = get_dates(self.trade_date)
+        self.assertListAlmostEqual([jp_yc.discount_factor(d) for d in dates],
+                                   self.yc.data[1:])
+        pickled_yc = loads(dumps(jp_yc))
+        self.assertListAlmostEqual([pickled_yc.discount_factor(d) for d in dates],
+                                   self.yc.data[1:])
+
     def test_contingentleg(self):
-        cl = ContingentLeg(datetime.date(2016, 11, 3),
-                           datetime.date(2021, 12, 20),
-                           12,
-                           True)
-        pickle = dumps(cl)
-        import pdb; pdb.set_trace()
-        cl2 = loads(pickle)
+        start_date = datetime.date(2016, 9, 20)
+        end_date = datetime.date(2021, 12, 20)
+        step_in_date = datetime.date(2016, 10, 7)
+        value_date = datetime.date(2016, 10, 11)
+        cl = ContingentLeg(start_date, end_date, 1)
+        fl = FeeLeg(start_date, end_date, True, 1, 1)
+        sc = SpreadCurve(self.trade_date, self.jp_yc, start_date,
+                         step_in_date, value_date, [end_date], array.array('d', [75*1e-4]), 0.4)
+        sc_pickled = loads(dumps(sc))
+        self.assertListAlmostEqual([sc.survival_probability(d) for d in get_dates(self.trade_date)],
+                                   [sc_pickled.survival_probability(d) for d in get_dates(self.trade_date)])
+        cl_pickled = loads(dumps(cl))
+        fl_pickled = loads(dumps(fl))
+        self.assertEqual(cl_pickled.pv(self.trade_date, step_in_date, value_date, self.jp_yc, sc, 0.4),
+                         cl.pv(self.trade_date, step_in_date, value_date, self.jp_yc, sc, 0.4))
+        self.assertEqual(fl.pv(self.trade_date, step_in_date, value_date, self.jp_yc, sc, False),
+                         fl_pickled.pv(self.trade_date, step_in_date, value_date, self.jp_yc, sc, False))