1 files changed, 66 insertions, 0 deletions
diff --git a/stator/tests/test_graph.py b/stator/tests/test_graph.py
new file mode 100644
index 0000000..f6b8404
--- /dev/null
+++ b/stator/tests/test_graph.py
@@ -0,0 +1,66 @@
+import pytest
+
+from stator.graph import State, StateGraph
+
+
+def test_declare():
+    """
+    Tests a basic graph declaration and various kinds of handler
+    lookups.
+    """
+
+    fake_handler = lambda: True
+
+    class TestGraph(StateGraph):
+        initial = State()
+        second = State()
+        third = State()
+        fourth = State()
+        final = State()
+
+        initial.add_transition(second, 60, handler=fake_handler)
+        second.add_transition(third, 60, handler="check_third")
+
+        def check_third(cls):
+            return True
+
+        @third.add_transition(fourth, 60)
+        def check_fourth(cls):
+            return True
+
+        fourth.add_manual_transition(final)
+
+    assert TestGraph.initial_state == TestGraph.initial
+    assert TestGraph.terminal_states == {TestGraph.final}
+
+    assert TestGraph.initial.children[TestGraph.second].get_handler() == fake_handler
+    assert (
+        TestGraph.second.children[TestGraph.third].get_handler()
+        == TestGraph.check_third
+    )
+    assert (
+        TestGraph.third.children[TestGraph.fourth].get_handler().__name__
+        == "check_fourth"
+    )
+
+
+def test_bad_declarations():
+    """
+    Tests that you can't declare an invalid graph.
+    """
+    # More than one initial state
+    with pytest.raises(ValueError):
+
+        class TestGraph(StateGraph):
+            initial = State()
+            initial2 = State()
+
+    # No initial states
+    with pytest.raises(ValueError):
+
+        class TestGraph(StateGraph):
+            loop = State()
+            loop2 = State()
+
+            loop.add_transition(loop2, 1, handler="fake")
+            loop2.add_transition(loop, 1, handler="fake")