class TestTakeTrainFunctionWithTimeBetweenStations(TestCase): """ Test to cover take train function. """ def setUp(self): """ (Line 2) 4 / \ 3 A1 / \ A -- B -- C -- D -- E -- F (Line 1) / \ | A2 2 | | A3 1 (Line 3) """ self.subway_system = SubwaySystem() self.subway_system.add_train_line(['A', 'B', 'C', 'D', 'E', 'F'], 'Line 1', [('A', 'B', 3), ('B', 'C', 5),('C', 'D', 5), ('D', 'E', 10), ('E', 'F', 15)]) self.subway_system.add_train_line(['1', '2', 'D', '3', '4'], 'Line 2', [('1', '2', 5), ('2', 'D', 5), ('D', '3', 1), ('3', '4', 10)]) self.subway_system.add_train_line(['A3', 'A2', 'F', 'A1', '4'], 'Line 3', [('A3', 'A2', 2), ('A2', 'F', 5), ('F', 'A1', 2), ('A1', '4', 10)]) def test_take_train_from_f_to_1(self): self.assertEqual(self.subway_system.take_train('F', '1'), (['F', 'A1', '4', '3', 'D', '2', '1'], 33)) def test_take_train_from_A_To_A3(self): """ Test take train function when going from A to A3 :return: """ self.assertEqual(self.subway_system.take_train('A', 'A3'), (['A', 'B', 'C', 'D', '3', '4', 'A1', 'F', 'A2', 'A3'], 43)) def test_take_train_from_A_To_A3_with_direct_connection(self): """ Test take train function when going from A to A3 Directly :return: """ self.subway_system.add_train_line(['A', 'A3'], "Subway to heaven", [('A', 'A3', 0.5)]) self.assertEqual(self.subway_system.take_train('A', 'A3'), (['A', 'A3'], 0.5)) def test_take_train_from_4_to_E(self): """ Test take train from 4 to E """ self.assertEqual(self.subway_system.take_train('4', 'E'), (['4', '3', 'D', 'E'], 21))
class TestTakeTrainFunctionWithNoTimeBetweenStations(TestCase): """ Test function take_train with no time is provided between train stations. """ def setUp(self): """ (Line 2) 4 / \ 3 A1 / \ A -- B -- C -- D -- E -- F (Line 1) / \ | A2 2 | | A3 1 (Line 3) """ self.subway_system = SubwaySystem() self.subway_system.add_train_line(['A', 'B', 'C', 'D', 'E', 'F'], 'Line 1') self.subway_system.add_train_line(['1', '2', 'D', '3', '4'], 'Line 2') self.subway_system.add_train_line(['A3', 'A2', 'F', 'A1', '4'], 'Line 3') def test_take_train_function_when_sending_A_A3(self): """ :return: """ self.assertEqual(self.subway_system.take_train("A", "A3"), (['A', 'B', 'C', 'D', 'E', 'F', 'A2', 'A3'], None)) def test_take_train_function_when_sending_A_A3_with_direct_route(self): """ Test take train function when having a direct route from A to A3 :return: """ self.subway_system.add_train_line(["A", "A3"], "Subway To Heaven", None) self.assertEqual(self.subway_system.take_train("A", "A3"), (['A', 'A3'], None)) def test_take_train_function_when_sending_A_A3_with_direct_route_backwards(self): """ Test take train function when having a direct route from A to A3 and going backwards :return: """ self.subway_system.add_train_line(["A", "A3"], "Subway To Heaven", None) self.assertEqual(self.subway_system.take_train("A3", "A"), (['A3', 'A'], None))