1 files changed, 120 insertions, 8 deletions

diff --git a/day9/test_init.py b/day9/test_init.py
index b36eed6..eb889c7 100644
--- a/day9/test_init.py
+++ b/day9/test_init.py
@@ -4,11 +4,123 @@ import os.path
 import day9
 
 
-def test_output_visualization():
-    expected = "..##..\n" "...##.\n" ".####.\n" "....#.\n" "s###.."
-
-    assignment = day9.AssignmentOne(path=os.path.dirname(day9.__file__))
-    unique_positions = assignment.unique_tail_positions(
-        input=assignment.read_input(True)
-    )
-    assert expected == assignment.visualize(6, 5, unique_positions)
+class TestAssignment:
+    def test_right(self):
+        expected = "\n".join(
+            [
+                "s###..",
+            ]
+        )
+
+        input = ["R 4"]
+
+        assignment = day9.AssignmentOne(path=os.path.dirname(day9.__file__))
+        unique_positions = assignment.unique_tail_positions(input=input, length=2)
+
+        assert assignment.visualize(range(6), range(1), unique_positions) == expected
+
+    def test_left(self):
+        expected = "\n".join(
+            [
+                "..###s",
+            ]
+        )
+
+        input = ["L 4"]
+
+        assignment = day9.AssignmentOne(path=os.path.dirname(day9.__file__))
+        unique_positions = assignment.unique_tail_positions(input=input, length=2)
+
+        assert (
+            assignment.visualize(range(-5, 1), range(1), unique_positions) == expected
+        )
+
+    def test_up(self):
+        expected = "\n".join(
+            [
+                ".",
+                ".",
+                "#",
+                "#",
+                "#",
+                "s",
+            ]
+        )
+
+        input = ["U 4"]
+
+        assignment = day9.AssignmentOne(path=os.path.dirname(day9.__file__))
+        unique_positions = assignment.unique_tail_positions(input=input, length=2)
+
+        assert assignment.visualize(range(1), range(6), unique_positions) == expected
+
+    def test_down(self):
+        expected = "\n".join(
+            [
+                "s",
+                "#",
+                "#",
+                "#",
+                ".",
+                ".",
+            ]
+        )
+
+        input = ["D 4"]
+
+        assignment = day9.AssignmentOne(path=os.path.dirname(day9.__file__))
+        unique_positions = assignment.unique_tail_positions(input=input, length=2)
+
+        assert (
+            assignment.visualize(range(1), range(-5, 1), unique_positions) == expected
+        )
+
+
+class TestAssignmentOne:
+    def test_output_visualization(self):
+        expected = "\n".join(["..##..", "...##.", ".####.", "....#.", "s###.."])
+
+        assignment = day9.AssignmentOne(path=os.path.dirname(day9.__file__))
+        unique_positions = assignment.unique_tail_positions(
+            input=assignment.read_input(True), length=2
+        )
+        assert assignment.visualize(range(6), range(5), unique_positions) == expected
+
+
+class TestAssignmentTwo:
+    def test_output_visualization(self):
+        expected = "\n".join(
+            [
+                "..........................",
+                "..........................",
+                "..........................",
+                "..........................",
+                "..........................",
+                "..........................",
+                "..........................",
+                "..........................",
+                "..........................",
+                "#.........................",
+                "#.............###.........",
+                "#............#...#........",
+                ".#..........#.....#.......",
+                "..#..........#.....#......",
+                "...#........#.......#.....",
+                "....#......s.........#....",
+                ".....#..............#.....",
+                "......#............#......",
+                ".......#..........#.......",
+                "........#........#........",
+                ".........########.........",
+            ]
+        )
+
+        assignment = day9.AssignmentTwo(path=os.path.dirname(day9.__file__))
+        unique_positions = assignment.unique_tail_positions(
+            input=assignment.read_input(True),
+            length=10,
+        )
+        assert (
+            assignment.visualize(range(-11, 15), range(-5, 16), unique_positions)
+            == expected
+        )