Adding callbacks

lib/fsmx.ex

@@ -2,19 +2,30 @@ defmodule Fsmx do
   @type state_t :: binary
 
   def transition(%mod{state: state} = struct, new_state) do
-    transitions = mod.__fsmx__(:transitions)
+    fsm = mod.__fsmx__()
+    transitions = fsm.__fsmx__(:transitions)
 
-    if valid_transition?(state, new_state, transitions) do
-      {:ok, %{struct | state: new_state}}
-    else
-      {:error, "invalid transition of #{mod} from #{state} to #{new_state}"}
+    with :ok <- validate_transition(state, new_state, transitions),
+         {:ok, struct} <- fsm.before_transition(struct, state, new_state),
+         {:ok, struct} <- do_transition(struct, new_state),
+         {:ok, struct} <- fsm.after_transition(struct, state, new_state) do
+      {:ok, struct}
     end
   end
 
-  defp valid_transition?(state, new_state, transitions) do
+  defp do_transition(struct, new_state) do
+    {:ok, %{struct | state: new_state}}
+  end
+
+  defp validate_transition(state, new_state, transitions) do
     transitions
     |> Map.get(state, [])
     |> is_or_contains?(new_state)
+    |> if do
+      :ok
+    else
+      {:error, "invalid transition from #{state} to #{new_state}"}
+    end
   end
 
   defp is_or_contains?(state, state), do: true

lib/fsmx/fsm.ex

@@ -0,0 +1,19 @@
+defmodule Fsmx.Fsm do
+  defmacro __using__(opts \\ []) do
+    quote do
+      @before_compile unquote(__MODULE__)
+
+      @fsm Keyword.get(unquote(opts), :fsm, __MODULE__)
+
+      def __fsmx__(:transitions), do: Keyword.fetch!(unquote(opts), :transitions)
+      def __fsmx__(:fsm), do: @fsm
+    end
+  end
+
+  defmacro __before_compile__(_env) do
+    quote generated: false do
+      def before_transition(struct, _from, _to), do: {:ok, struct}
+      def after_transition(struct, _from, _to), do: {:ok, struct}
+    end
+  end
+end

lib/fsmx/struct.ex

@@ -1,7 +1,13 @@
 defmodule Fsmx.Struct do
-  defmacro __using__(opts) do
+  defmacro __using__(opts \\ []) do
     quote do
-      def __fsmx__(:transitions), do: Keyword.fetch!(unquote(opts), :transitions)
+      @fsm Keyword.get(unquote(opts), :fsm, __MODULE__)
+
+      if @fsm == __MODULE__ do
+        use Fsmx.Fsm, unquote(opts)
+      end
+
+      def __fsmx__, do: @fsm
     end
   end
 end

test/fsmx/struct_test.exs

@@ -0,0 +1,5 @@
+defmodule Fsmx.StructTest do
+  use ExUnit.Case
+
+  alias Fsmx.Struct
+end

test/fsmx_test.exs

@@ -2,7 +2,7 @@ defmodule FsmxTest do
   use ExUnit.Case
   doctest Fsmx
 
-  alias Fsmx.TestStructs.{Simple, WithCallbacks}
+  alias Fsmx.TestStructs.{Simple, WithCallbacks, WithSeparateFsm}
 
   describe "transition/2" do
     test "can do simple transitions" do
@@ -22,12 +22,49 @@ defmodule FsmxTest do
 
       assert {:error, msg} = Fsmx.transition(one, "3")
 
-      assert msg == "invalid transition of Elixir.Fsmx.TestStructs.Simple from 1 to 3"
+      assert msg == "invalid transition from 1 to 3"
     end
   end
 
-  describe "transition/2 with callbacks" do
+  describe "transition/2 with before_callbacks" do
     test "calls before_transition/2 on struct" do
+      one = %WithCallbacks.ValidBefore{state: "1", before: false, after: false}
+
+      {:ok, two} = Fsmx.transition(one, "2")
+
+      assert %WithCallbacks.ValidBefore{state: "2", before: true, after: false} = two
+    end
+
+    test "fails to transition if before_transition/3 returns an error" do
+      one = %WithCallbacks.InvalidBefore{state: "1", before: false, after: false}
+
+      {:error, :before_failed} = Fsmx.transition(one, "2")
+    end
+  end
+
+  describe "transition/2 with after_callbacks" do
+    test "calls before_transition/2 on struct" do
+      one = %WithCallbacks.ValidAfter{state: "1", before: false, after: false}
+
+      {:ok, two} = Fsmx.transition(one, "2")
+
+      assert %WithCallbacks.ValidAfter{state: "2", before: false, after: true} = two
+    end
+
+    test "fails to transition if before_transition/3 returns an error" do
+      one = %WithCallbacks.InvalidAfter{state: "1", before: false, after: false}
+
+      {:error, :after_failed} = Fsmx.transition(one, "2")
+    end
+  end
+
+  describe "transition/2 with separate fsm module" do
+    test "works just the same" do
+      one = %WithSeparateFsm{state: "1"}
+
+      {:ok, two} = Fsmx.transition(one, "2")
+
+      assert %WithSeparateFsm{state: "2", before: true, after: true} = two
     end
   end
 end

test/support/test_structs/with_callbacks.ex

@@ -1,9 +1,51 @@
 defmodule Fsmx.TestStructs.WithCallbacks do
-  defstruct [:state]
+  defmodule ValidBefore do
+    defstruct state: "1", before: false, after: false
 
-  use Fsmx.Struct, transitions: %{"1" => ["2"], "2" => ["3"]}
+    use Fsmx.Struct, transitions: %{"1" => "2"}
 
-  def before_transition(struct, state) do
-    IO.inspect(struct, state)
+    def before_transition(struct, _old_state, _new_state) do
+      {:ok, %{struct | before: true}}
+    end
+  end
+
+  defmodule ValidAfter do
+    defstruct state: "1", before: false, after: false
+
+    use Fsmx.Struct, transitions: %{"1" => "2"}
+
+    def after_transition(struct, _old_state, _new_state) do
+      {:ok, %{struct | after: true}}
+    end
+  end
+
+  defmodule InvalidBefore do
+    defstruct state: "1", before: false, after: false
+
+    use Fsmx.Struct, transitions: %{"1" => "2"}
+
+    def before_transition(struct, _old_state, _new_state) do
+      {:error, :before_failed}
+    end
+  end
+
+  defmodule InvalidAfter do
+    defstruct state: "1", before: false, after: false
+
+    use Fsmx.Struct, transitions: %{"1" => "2"}
+
+    def after_transition(struct, _old_state, _new_state) do
+      {:error, :after_failed}
+    end
+  end
+
+  defmodule PartialCallback do
+    defstruct state: "1", before: false, after: false
+
+    use Fsmx.Struct, transitions: %{"1" => "2", "2" => "3"}
+
+    def after_transition(struct, "1", "2") do
+      {:ok, struct}
+    end
   end
 end

test/support/test_structs/with_separate_fsm.ex

@@ -0,0 +1,13 @@
+defmodule Fsmx.TestStructs.WithSeparateFsm do
+  defstruct [:state, before: false, after: false]
+
+  use Fsmx.Struct, fsm: __MODULE__.Fsm
+
+  defmodule Fsm do
+    use Fsmx.Fsm, transitions: %{"1" => "2"}
+
+    def before_transition(struct, _, _), do: {:ok, %{struct | before: true}}
+
+    def after_transition(struct, _, _), do: {:ok, %{struct | after: true}}
+  end
+end