day09: tests

day09/day9.py

@@ -3,6 +3,9 @@
 from dataclasses import dataclass
 from typing import Callable
 
+INPUT = "day09/input.txt"
+INPUT_SMALL = "day09/input-small.txt"
+
 
 @dataclass
 class ValueArray:
@@ -53,10 +56,10 @@ def calculate_value(array: list[int], array_below: list[int]) -> int:
         self.generic_extrapolate(add_to_array, calculate_value)
 
 
-def get_input() -> list[ValueArray]:
+def get_input(path: str) -> list[ValueArray]:
     """turns inputs into nice ValueArrays"""
     result = []
-    with open("day09/input.txt", "r", encoding="utf8") as file:
+    with open(path, "r", encoding="utf8") as file:
         for line in file:
             values = ValueArray([[int(item) for item in line.split()]])
             result.append(values)
@@ -75,20 +78,28 @@ def interpolate(values: list[int]) -> list[int]:
     return result
 
 
-def main() -> None:
-    """main function"""
-    inputs = get_input()
-
-    # q1
+def part1(inputs: list[ValueArray]) -> int:
     for input in inputs:
         input.extrapolate_right()
+    return sum(input.sub_arrays[0][-1] for input in inputs)
 
-    print(sum(input.sub_arrays[0][-1] for input in inputs))
-    # q2
+
+def part2(inputs: list[ValueArray]) -> int:
     for input in inputs:
         input.extrapolate_left()
 
-    print(sum(input.sub_arrays[0][0] for input in inputs))
+    return sum(input.sub_arrays[0][0] for input in inputs)
+
+
+def main() -> None:
+    """main function"""
+    inputs = get_input(INPUT)
+
+    # q1
+    print(part1(inputs))
+
+    # q2
+    print(part2(inputs))
 
 
 if __name__ == "__main__":

day09/input-small.txt

@@ -0,0 +1,3 @@
+0 3 6 9 12 15
+1 3 6 10 15 21
+10 13 16 21 30 45

day09/tests/test_day9.py

@@ -0,0 +1,23 @@
+from day09.day9 import INPUT_SMALL, ValueArray, get_input, interpolate, part1, part2
+
+
+def test_get_input() -> None:
+    values: list[ValueArray] = get_input(INPUT_SMALL)
+    assert len(values) == 3
+    assert len(values[0].sub_arrays) == 3
+    assert values[0].sub_arrays[0] == [0, 3, 6, 9, 12, 15]
+
+
+def test_interpolate() -> None:
+    values = [0, 3, 6, 9, 12, 15]
+    assert interpolate(values) == [3, 3, 3, 3, 3]
+
+
+def test_part1() -> None:
+    values: list[ValueArray] = get_input(INPUT_SMALL)
+    assert part1(values) == 114
+
+
+def test_part2() -> None:
+    values: list[ValueArray] = get_input(INPUT_SMALL)
+    assert part2(values) == 2