#!/usr/bin/python

import unittest

import main



class ExamplesPart1(unittest.TestCase):
    def test_example1(self):
        data = ">"
        self.assertEqual(main.solve_part1(data), 2)

    def test_example2(self):
        data = "^>v<"
        self.assertEqual(main.solve_part1(data), 4)

    def test_example3(self):
        data = "^v^v^v^v^v"
        self.assertEqual(main.solve_part1(data), 2)



class ExamplesPart2(unittest.TestCase):
    def test_example1(self):
        data = "^v"
        self.assertEqual(main.solve_part2(data), 3)

    def test_example2(self):
        data = "^>v<"
        self.assertEqual(main.solve_part2(data), 3)

    def test_example3(self):
        data = "^v^v^v^v^v"
        self.assertEqual(main.solve_part2(data), 11)



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