import unittest
from common import OASISLine


class Day9Tests(unittest.TestCase):
    def test_part1(self):
        analysis = OASISLine.parse("example_input.txt")
        extrapolation = map(lambda o: o.reduce().extrapolate(), analysis)
        self.assertEqual(sum(extrapolation), 114)

    def test_part2(self):
        analysis = OASISLine.parse("example_input.txt")
        extrapolation = map(lambda o: o.reduce().extrapolate_back(), analysis)
        self.assertEqual(sum(extrapolation), 2)

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