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))