tests do not rely on private code

1 files changed, 15 insertions, 26 deletions

diff --git a/tests/test_pickle.py b/tests/test_pickle.py
index 476e834..4bb1d46 100644
--- a/tests/test_pickle.py
+++ b/tests/test_pickle.py
@@ -4,19 +4,11 @@ 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
+from pyisda.utils import build_yc
 
 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)
+    yc = build_yc(trade_date, ql_curve = True)
 
     def assertListAlmostEqual(self, l1, l2):
         if len(l1) != len(l2):
@@ -25,29 +17,26 @@ class TestPickle(unittest.TestCase):
             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:])
+        orig_dfs = [self.yc.discount_factor(d) for d in self.yc.dates]
+        pickled_yc = loads(dumps(self.yc))
+        self.assertListAlmostEqual([pickled_yc.discount_factor(d) for d in pickled_yc.dates],
+                                   orig_dfs)
 
-    def test_contingentleg(self):
-        start_date = datetime.date(2016, 9, 20)
+    def test_legs(self):
+        start_date = datetime.date(2016, 3, 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,
+        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))
-        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)])
+        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))
-        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))
+        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))