From f01817f2bc26e3bca880e0c2e817b6414f198c20 Mon Sep 17 00:00:00 2001 From: Pietro Date: Sat, 10 Apr 2021 14:01:05 +0200 Subject: [PATCH] Added Tests for TrajectoryGenerator --- test_trajectorygenerator.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/test_trajectorygenerator.py b/test_trajectorygenerator.py index a35c513..b0535c2 100644 --- a/test_trajectorygenerator.py +++ b/test_trajectorygenerator.py @@ -22,12 +22,17 @@ class TestTrajectoryGenerator(unittest.TestCase): tg = TrajectoryGenerator(self.j1) end_time = random.randint(5, 100) sigma = tg.CTBN_Sample(end_time) - self.assertLessEqual(sigma.times.loc[len(sigma) - 1].at["Time"], end_time) - for index, row in sigma.iterrows(): + self.assertLessEqual(sigma.times[len(sigma.times) - 1], end_time) + for index in range(len(sigma.times)): if index > 0: - self.assertLess(sigma.times.loc[index - 1].at["Time"], row.at["Time"]) - diff = abs(sum(sigma.loc[index - 1, sigma.columns != "Time"]) - - sum(sigma.loc[index, sigma.columns != "Time"])) + self.assertLess(sigma.times[index - 1], sigma.times[index]) + diff = abs(sum(sigma.trajectory[index - 1]) - sum(sigma.trajectory[index])) self.assertEqual(diff, 1) + def test_generated_trajectory_max_tr(self): + tg = TrajectoryGenerator(self.j1) + n_tr = random.randint(5, 100) + sigma = tg.CTBN_Sample(max_tr = n_tr) + self.assertEqual(len(sigma.times), n_tr + 1) + unittest.main() \ No newline at end of file