diff --git a/src/lib.rs b/src/lib.rs
index e0d6cab..41c433a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -33,13 +33,15 @@ pub trait GraphLike<Node, NodeId> {
     {
         self.iter_ids().filter(|id| self.get_node(*id).is_none())
     }
-    fn iter_neighbour_nodes(&self, id: NodeId) -> Vec<(NodeId, Option<&Node>)>
+    fn iter_neighbour_nodes<'a>(
+        &'a self,
+        id: NodeId,
+    ) -> impl Iterator<Item = (NodeId, Option<&'a Node>)>
     where
         NodeId: Copy,
+        Node: 'a,
     {
-        self.iter_neighbor_ids(id)
-            .map(|id| (id, self.get_node(id)))
-            .collect()
+        self.iter_neighbor_ids(id).map(|id| (id, self.get_node(id)))
     }
 
     fn single_step(
@@ -130,3 +132,297 @@ pub trait Ruleset<Node, NodeId> {
             .map(|(id, _)| id)
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashSet;
+
+    use crate::{GraphLike, Ruleset};
+
+    struct TestGraph {
+        tiles: [[Option<u8>; 3]; 3],
+    }
+
+    impl GraphLike<u8, (usize, usize)> for TestGraph {
+        fn get_node(&self, id: (usize, usize)) -> Option<&u8> {
+            self.tiles[id.0][id.1].as_ref()
+        }
+
+        fn set_node(&mut self, id: (usize, usize), node: Option<u8>) {
+            self.tiles[id.0][id.1] = node;
+        }
+
+        fn iter_ids(&self) -> impl Iterator<Item = (usize, usize)> {
+            (0..3).flat_map(|i| (0..3).map(move |j| (i, j)))
+        }
+
+        fn iter_neighbor_ids(&self, id: (usize, usize)) -> impl Iterator<Item = (usize, usize)> {
+            let mut neighbors = vec![];
+            if id.0 > 0 {
+                neighbors.push((id.0 - 1, id.1));
+            }
+            if id.0 < 2 {
+                neighbors.push((id.0 + 1, id.1));
+            }
+            if id.1 > 0 {
+                neighbors.push((id.0, id.1 - 1));
+            }
+            if id.1 < 2 {
+                neighbors.push((id.0, id.1 + 1));
+            }
+            neighbors.into_iter()
+        }
+    }
+
+    /// Entropy is the number of empty neighbors, and choose sets the value to the same value.
+    /// Just for testing.
+    struct TestRuleset;
+
+    impl Ruleset<u8, (usize, usize)> for TestRuleset {
+        fn entropy(
+            &self,
+            graph: &impl GraphLike<u8, (usize, usize)>,
+            id: (usize, usize),
+        ) -> Option<f32> {
+            let neighbors = graph
+                .iter_neighbour_nodes(id)
+                .filter(|n| n.1.is_some())
+                .count();
+            Some(4_usize.strict_sub(neighbors) as f32)
+        }
+
+        fn choose(
+            &self,
+            graph: &impl GraphLike<u8, (usize, usize)>,
+            id: (usize, usize),
+            _retry_count: Option<usize>,
+        ) -> Option<u8> {
+            let neighbors = graph
+                .iter_neighbour_nodes(id)
+                .filter(|n| n.1.is_some())
+                .count();
+            Some(4_usize.strict_sub(neighbors) as u8)
+        }
+    }
+
+    #[test]
+    fn get_node_works() {
+        // This really tests that the implementation itself works...
+        let graph = TestGraph {
+            tiles: [
+                [None, None, None],
+                [None, None, Some(99)],
+                [None, None, None],
+            ],
+        };
+        assert_eq!(graph.get_node((0, 0)), None);
+        assert_eq!(graph.get_node((1, 2)), Some(&99));
+    }
+
+    #[test]
+    fn set_node_works() {
+        // This really tests that the implementation itself works...
+        let mut graph = TestGraph {
+            tiles: [[None; 3]; 3],
+        };
+        graph.set_node((1, 2), Some(99));
+        assert_eq!(
+            graph.tiles,
+            [
+                [None, None, None],
+                [None, None, Some(99)],
+                [None, None, None]
+            ]
+        );
+    }
+
+    #[test]
+    fn iter_ids_works() {
+        // This really tests that the implementation itself works...
+        let graph = TestGraph {
+            tiles: [[None; 3]; 3],
+        };
+        assert_eq!(
+            graph.iter_ids().collect::<HashSet<(usize, usize)>>(),
+            vec![
+                (0, 0),
+                (0, 1),
+                (0, 2),
+                (1, 0),
+                (1, 1),
+                (1, 2),
+                (2, 0),
+                (2, 1),
+                (2, 2)
+            ]
+            .into_iter()
+            .collect::<HashSet<(usize, usize)>>()
+        );
+    }
+
+    #[test]
+    fn iter_nodes_works() {
+        let graph = TestGraph {
+            tiles: [
+                [None, None, None],
+                [None, None, Some(99)],
+                [None, None, None],
+            ],
+        };
+        assert_eq!(
+            graph
+                .iter_nodes()
+                .collect::<HashSet<((usize, usize), Option<&u8>)>>(),
+            vec![
+                ((0, 0), None),
+                ((0, 1), None),
+                ((0, 2), None),
+                ((1, 0), None),
+                ((1, 1), None),
+                ((1, 2), Some(&99)),
+                ((2, 0), None),
+                ((2, 1), None),
+                ((2, 2), None)
+            ]
+            .into_iter()
+            .collect::<HashSet<((usize, usize), Option<&u8>)>>()
+        );
+    }
+
+    #[test]
+    fn iter_neighbor_ids_works() {
+        // This really tests that the implementation itself works...
+        let graph = TestGraph {
+            tiles: [[None; 3]; 3],
+        };
+        assert_eq!(
+            graph
+                .iter_neighbor_ids((1, 1))
+                .collect::<HashSet<(usize, usize)>>(),
+            vec![(0, 1), (2, 1), (1, 0), (1, 2)]
+                .into_iter()
+                .collect::<HashSet<(usize, usize)>>()
+        );
+    }
+
+    #[test]
+    fn iter_neighbour_nodes_works() {
+        let graph = TestGraph {
+            tiles: [
+                [None, None, None],
+                [None, None, Some(99)],
+                [None, None, None],
+            ],
+        };
+        assert_eq!(
+            graph
+                .iter_neighbour_nodes((1, 1))
+                .collect::<HashSet<((usize, usize), Option<&u8>)>>(),
+            vec![
+                ((0, 1), None),
+                ((2, 1), None),
+                ((1, 0), None),
+                ((1, 2), Some(&99))
+            ]
+            .into_iter()
+            .collect::<HashSet<((usize, usize), Option<&u8>)>>()
+        );
+    }
+
+    #[test]
+    fn entropy_works() {
+        // This really tests that the implementation itself works...
+        let graph = TestGraph {
+            tiles: [
+                [None, None, None],
+                [None, None, None],
+                [None, None, Some(99)],
+            ],
+        };
+        let ruleset = TestRuleset;
+        assert_eq!(ruleset.entropy(&graph, (1, 1)), Some(4.0));
+        assert_eq!(ruleset.entropy(&graph, (0, 0)), Some(4.0));
+        assert_eq!(ruleset.entropy(&graph, (2, 1)), Some(3.0));
+    }
+
+    #[test]
+    fn choose_works() {
+        // This really tests that the implementation itself works...
+        let graph = TestGraph {
+            tiles: [
+                [None, None, None],
+                [None, None, Some(0)],
+                [None, Some(0), None],
+            ],
+        };
+        let ruleset = TestRuleset;
+        assert_eq!(ruleset.entropy(&graph, (0, 0)), Some(4.0));
+        assert_eq!(ruleset.entropy(&graph, (1, 1)), Some(2.0));
+        assert_eq!(ruleset.entropy(&graph, (2, 2)), Some(2.0));
+    }
+
+    #[test]
+    fn find_lowest_entropy_works() {
+        let graph = TestGraph {
+            tiles: [
+                [None, None, None],
+                [Some(0), None, Some(0)],
+                [None, Some(0), None],
+            ],
+        };
+        let ruleset = TestRuleset;
+        assert_eq!(ruleset.find_lowest_entropy(&graph, None), Some((1, 1)));
+    }
+
+    #[test]
+    fn single_step_works() {
+        let mut graph = TestGraph {
+            tiles: [
+                [None, None, None],
+                [Some(0), None, Some(0)],
+                [None, Some(0), Some(0)],
+            ],
+        };
+        let ruleset = TestRuleset;
+        assert_eq!(graph.single_step(&ruleset, None), Some((1, 1)));
+        assert_eq!(
+            graph.tiles,
+            [
+                [None, None, None],
+                [Some(0), Some(1), Some(0)],
+                [None, Some(0), Some(0)],
+            ]
+        );
+        assert_eq!(graph.single_step(&ruleset, None), Some((2, 0)));
+        assert_eq!(
+            graph.tiles,
+            [
+                [None, None, None],
+                [Some(0), Some(1), Some(0)],
+                [Some(2), Some(0), Some(0)],
+            ]
+        );
+    }
+
+    #[test]
+    fn full_generation_works() {
+        let mut graph = TestGraph {
+            // tiles: [[None; 3]; 3],
+            tiles: [
+                [None, Some(99), None],
+                [Some(99), None, Some(99)],
+                [None, Some(99), None],
+            ],
+        };
+        let ruleset = TestRuleset;
+        assert_eq!(graph.full_generation(&ruleset), Ok(()));
+        assert_eq!(
+            graph.tiles,
+            [
+                [Some(2), Some(99), Some(2)],
+                [Some(99), Some(0), Some(99)],
+                [Some(2), Some(99), Some(2)],
+            ]
+        );
+    }
+}