class TestGraph(unittest.TestCase):
    def test_add_node(self):
        # given
        graph = Graph()
        graph2 = Graph()
        graph3 = Graph()

        # when
        graph.add_node('A')
        graph.add_node('B')
        graph.add_node('C')
        graph.add_node('D')

        graph2.add_node(1)
        graph2.add_node(2)
        graph2.add_node(3)
        graph2.add_node(4)

        # then
        self.assertEqual('A' in graph.nodes, True)
        self.assertEqual('B' in graph.nodes, True)
        self.assertEqual('C' in graph.nodes, True)
        self.assertEqual('D' in graph.nodes, True)
        self.assertEqual(1 in graph2.nodes, True)
        self.assertEqual(2 in graph2.nodes, True)
        self.assertEqual(3 in graph2.nodes, True)
        self.assertEqual(4 in graph2.nodes, True)

    def test_add_edges(self):
        # given
        graph = Graph()

        # when
        graph.add_edge('A', 'B', 20)
        graph.add_edge('A', 'C', 10)
        graph.add_edge('B', 'D', 15)
        graph.add_edge('C', 'D', 10)
        graph.add_edge('B', 'E', 50)
        graph.add_edge('D', 'E', 30)
        graph.add_edge('E', 'F', 5)
        graph.add_edge('F', 'G', 2)

        # then
        self.assertEqual(graph.edges['A'], ['B', 'C'])
        self.assertEqual(graph.edges['B'], ['A', 'D', 'E'])
        self.assertEqual(graph.edges['C'], ['A', 'D'])
        self.assertEqual(graph.edges['D'], ['B', 'C', 'E'])
        self.assertEqual(graph.edges['E'], ['B', 'D', 'F'])
        self.assertEqual(graph.edges['F'], ['E', 'G'])
        self.assertEqual(graph.edges['G'], ['F'])

        self.assertEqual(graph.distances[('A', 'B')], 20)
        self.assertEqual(graph.distances[('A', 'C')], 10)
        self.assertEqual(graph.distances[('B', 'D')], 15)
        self.assertEqual(graph.distances[('C', 'D')], 10)
        self.assertEqual(graph.distances[('B', 'E')], 50)
        self.assertEqual(graph.distances[('D', 'E')], 30)
        self.assertEqual(graph.distances[('E', 'F')], 5)
        self.assertEqual(graph.distances[('F', 'G')], 2)
        self.assertEqual(graph.distances[('B', 'A')], 20)
        self.assertEqual(graph.distances[('C', 'A')], 10)
        self.assertEqual(graph.distances[('D', 'B')], 15)
        self.assertEqual(graph.distances[('D', 'C')], 10)
        self.assertEqual(graph.distances[('E', 'B')], 50)
        self.assertEqual(graph.distances[('E', 'D')], 30)
        self.assertEqual(graph.distances[('F', 'E')], 5)
        self.assertEqual(graph.distances[('G', 'F')], 2)

    def test_shortest_path(self):
        # given
        graph = Graph()

        # when
        graph.add_edge('A', 'B', 20)
        graph.add_edge('A', 'C', 10)
        graph.add_edge('B', 'D', 15)
        graph.add_edge('C', 'D', 10)
        graph.add_edge('B', 'E', 50)
        graph.add_edge('D', 'E', 30)
        graph.add_edge('E', 'F', 5)
        graph.add_edge('F', 'G', 2)

        #then
        self.assertEqual(shortest_path(graph, 'A', 'D'), (20, ['A', 'C', 'D']))
        self.assertEqual(shortest_path(graph, 'D', 'A'), (20, ['D', 'C', 'A']))
        self.assertEqual(shortest_path(graph, 'A', 'G'), (57, ['A', 'C', 'D', 'E', 'F', 'G']))
        self.assertEqual(shortest_path(graph, 'G', 'A'), (57, ['G', 'F', 'E', 'D', 'C', 'A']))
        self.assertEqual(shortest_path(graph, 'C', 'G'), (47, ['C', 'D', 'E', 'F', 'G']))
        self.assertEqual(shortest_path(graph, 'G', 'C'), (47, ['G', 'F', 'E', 'D', 'C']))
        self.assertEqual(shortest_path(graph, 'E', 'B'), (45, ['E', 'D', 'B']))
        self.assertEqual(shortest_path(graph, 'B', 'E'), (45, ['B', 'D', 'E']))


if __name__ == '__main__':
   unittest.main()