From 0a3c07f2cf2ed8f6bcca899859eba57322ebb1b6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Lu=C3=ADs=20Ferreira?= <luis@zamith.pt>
Date: Wed, 9 Aug 2023 08:41:40 +0100
Subject: [PATCH] Allows for multiple states machines in the same struct (#16)

* Allows for multiple states machines in the same struct

Why:

* In certain situations you might want to allow for multiple state
  machines in the same struct.

  See #13

This change addresses the need by:

* Allowing for Fsmx.Struct to be included multiple times per module,
  with a different `state_field`, so it runs independently. When calling
  any of the transition function you can optionally pass it the field
  name, so you can control which one to transition. All of the other
  options, including `transitions` are also independent
* Passing the `state_field` to the callbacks as well, so that you can
  have them be different depending on the which field is being
  transitionned. As of now, the old version is still supported, but it
  is encouraged that you use the new version, which will make it simpler
  to add new fields if you wish.

* Bump the elixir versions we test against

* Use OTP version 25

* Adds the new functionality to the README

* Bump the min elixir version supported and the minor version

* Corrects typespecs and uses `state_field` everywhere
---
 .github/workflows/test.yml                    |  4 +-
 README.md                                     | 41 ++++++++-
 config/config.exs                             |  4 +-
 config/dev.exs                                |  2 +-
 config/prod.exs                               |  2 +-
 config/test.exs                               |  2 +-
 lib/fsmx.ex                                   | 83 +++++++++++++------
 lib/fsmx/fsm.ex                               | 42 +++++++---
 lib/fsmx/struct.ex                            | 33 +++++++-
 mix.exs                                       |  4 +-
 mix.lock                                      |  8 +-
 .../20200624134635_create_test_schemas.exs    |  1 +
 test/fsmx/ecto_test.exs                       | 38 ++++++++-
 test/fsmx/struct_test.exs                     | 38 ++++++++-
 test/support/test_ecto_schemas/multi_state.ex | 12 +++
 .../test_ecto_schemas/with_callbacks.ex       | 21 +++++
 test/support/test_structs/multi_state.ex      | 10 +++
 17 files changed, 285 insertions(+), 60 deletions(-)
 create mode 100644 test/support/test_ecto_schemas/multi_state.ex
 create mode 100644 test/support/test_structs/multi_state.ex

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 03afd86..7268ddf 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -14,11 +14,11 @@ jobs:
   test:
     strategy:
       matrix:
-        elixir: ['1.8', '1.9', '1.10']
+        elixir: ['1.13', '1.14', '1.15']
 
     runs-on: ubuntu-latest
     container:
-      image: elixir:${{ matrix.elixir }}-alpine
+      image: elixir:${{ matrix.elixir }}-otp-25-alpine
       env:
         POSTGRES_HOST: postgres
         MIX_ENV: test
diff --git a/README.md b/README.md
index f582a27..79ffb0d 100644
--- a/README.md
+++ b/README.md
@@ -51,7 +51,7 @@ defmodule App.StateMachine do
     "two" => ["three", "four"],
     "three" => "four",
     "four" => :*, # can transition to any state
-    "*" => ["five"] # can transition from any state to "five"
+    :* => ["five"] # can transition from any state to "five"
   }
 end
 ```
@@ -148,6 +148,45 @@ defmodule App.BusinessLogic do
 end
 ```
 
+### Multiple state machines in the same struct
+
+Not all structs have a single state machine, sometimes you might need more,
+using different fields for that effect. Here's how you can do it:
+
+```elixir
+defmodule App.StateMachine do
+  defstruct [:state, :other_state, :data]
+
+  use Fsmx.Struct, transitions: %{
+    "one" => ["two", "three"],
+    "two" => ["three", "four"],
+    "three" => "four",
+    "four" => :*, # can transition to any state
+    :* => ["five"] # can transition from any state to "five"
+  }
+
+  use Fsmx.Struct,
+    state_field: :other_state,
+    transitions: %{
+        "initial" => ["middle", "middle2"],
+        "middle" => "middle2",
+        :* => "final"
+    }
+end
+```
+
+Use it via the `Fsmx.transition/3` function:
+
+```elixir
+struct = %App.StateMachine{state: "one", other_state: "initial", data: nil}
+
+Fsmx.transition(struct, "two")
+# {:ok, %App.StateMachine{state: "two", other_state: "initial"}}
+
+Fsmx.transition(struct, "final", field: :other_state)
+# {:ok, %App.StateMachine{state: "one", other_state: "final"}}
+```
+
 ## Ecto support
 
 Support for Ecto is built in, as long as `ecto` is in your `mix.exs` dependencies. With it, you get the ability to
diff --git a/config/config.exs b/config/config.exs
index 8233fe9..d1186fe 100644
--- a/config/config.exs
+++ b/config/config.exs
@@ -1,3 +1,3 @@
-use Mix.Config
+import Config
 
-import_config "#{Mix.env()}.exs"
+import_config "#{config_env()}.exs"
diff --git a/config/dev.exs b/config/dev.exs
index d2d855e..becde76 100644
--- a/config/dev.exs
+++ b/config/dev.exs
@@ -1 +1 @@
-use Mix.Config
+import Config
diff --git a/config/prod.exs b/config/prod.exs
index d2d855e..becde76 100644
--- a/config/prod.exs
+++ b/config/prod.exs
@@ -1 +1 @@
-use Mix.Config
+import Config
diff --git a/config/test.exs b/config/test.exs
index 345eed9..be4e402 100644
--- a/config/test.exs
+++ b/config/test.exs
@@ -1,4 +1,4 @@
-use Mix.Config
+import Config
 
 config :fsmx, ecto_repos: [Fsmx.Repo]
 
diff --git a/lib/fsmx.ex b/lib/fsmx.ex
index f3ab891..f4e7fba 100644
--- a/lib/fsmx.ex
+++ b/lib/fsmx.ex
@@ -2,29 +2,40 @@ defmodule Fsmx do
   @moduledoc """
   """
 
+  alias Fsmx.Fsm
+
   @type state_t :: binary() | atom()
+  @type opts_t :: [state_field: atom()]
+
+  @spec transition(struct(), state_t(), opts_t()) :: {:ok, struct} | {:error, any}
+  def transition(%{} = struct, new_state, opts \\ []) do
+    opts = Keyword.put_new(opts, :state_field, Fsm.default_state_field())
+    state_field = Keyword.get(opts, :state_field)
 
-  @spec transition(struct(), state_t()) :: {:ok, struct} | {:error, any}
-  def transition(%mod{} = struct, new_state) do
-    fsm = mod.__fsmx__()
-    with {:ok, struct} <- before_transition(struct, new_state) do
-      state_field = fsm.__fsmx__(:state_field)
-      {:ok, struct |> Map.put(state_field,  new_state)}
+    with {:ok, struct} <- before_transition(struct, new_state, state_field) do
+      {:ok, struct |> Map.put(state_field, new_state)}
     end
   end
 
   if Code.ensure_loaded?(Ecto) do
-    @spec transition_changeset(struct(), state_t(), map) :: Ecto.Changeset.t()
-    def transition_changeset(%mod{} = schema, new_state, params \\ %{}) do
-      fsm = mod.__fsmx__()
-      state_field = fsm.__fsmx__(:state_field)
+    @spec transition_changeset(struct(), state_t(), map, opts_t()) :: Ecto.Changeset.t()
+    def transition_changeset(%mod{} = schema, new_state, params \\ %{}, opts \\ []) do
+      opts = Keyword.put_new(opts, :state_field, Fsm.default_state_field())
+      state_field = Keyword.get(opts, :state_field)
       state = schema |> Map.fetch!(state_field)
+      fsm = mod.__fsmx__(state_field)
 
-      with {:ok, schema} <- before_transition(schema, new_state) do
+      with {:ok, schema} <- before_transition(schema, new_state, state_field) do
         schema
         |> Ecto.Changeset.change()
         |> Ecto.Changeset.put_change(state_field, new_state)
-        |> fsm.transition_changeset(state, new_state, params)
+        |> then(fn changeset ->
+          if function_exported?(fsm, :transition_changeset, 4) do
+            fsm.transition_changeset(changeset, state, new_state, params)
+          else
+            fsm.transition_changeset(changeset, state, new_state, params, state_field)
+          end
+        end)
       else
         {:error, msg} ->
           schema
@@ -33,39 +44,54 @@ defmodule Fsmx do
       end
     end
 
-    @spec transition_multi(Ecto.Multi.t(), struct(), any, state_t, map) :: Ecto.Multi.t()
-    def transition_multi(multi, %mod{} = schema, id, new_state, params \\ %{}) do
-      fsm = mod.__fsmx__()
-      state = schema |> Map.fetch!(fsm.__fsmx__(:state_field))
+    @spec transition_multi(Ecto.Multi.t(), struct(), any, state_t, map, opts_t()) ::
+            Ecto.Multi.t()
+    def transition_multi(multi, %mod{} = schema, id, new_state, params \\ %{}, opts \\ []) do
+      opts = Keyword.put_new(opts, :state_field, Fsm.default_state_field())
+      state_field = Keyword.get(opts, :state_field)
+      state = schema |> Map.fetch!(state_field)
+      fsm = mod.__fsmx__(state_field)
 
-      changeset = transition_changeset(schema, new_state, params)
+      changeset = transition_changeset(schema, new_state, params, state_field: state_field)
 
       multi
       |> Ecto.Multi.update(id, changeset)
       |> Ecto.Multi.run("#{id}-callback", fn _repo, changes ->
-        fsm.after_transition_multi(Map.fetch!(changes, id), state, new_state)
+        if function_exported?(fsm, :after_transition_multi, 3) do
+          fsm.after_transition_multi(Map.fetch!(changes, id), state, new_state)
+        else
+          fsm.after_transition_multi(Map.fetch!(changes, id), state, new_state, state_field)
+        end
       end)
     end
   end
 
-  defp before_transition(%mod{} = struct, new_state) do
-    fsm = mod.__fsmx__()
-    state = struct |> Map.fetch!(fsm.__fsmx__(:state_field))
-    transitions = fsm.__fsmx__(:transitions)
+  defp before_transition(%mod{} = struct, new_state, state_field) do
+    fsm = mod.__fsmx__(state_field)
+    state = struct |> Map.fetch!(state_field)
+    transitions = fsm.__fsmx__(state_field, :transitions)
 
-    with :ok <- validate_transition(state, new_state, transitions) do
-      fsm.before_transition(struct, state, new_state)
+    with :ok <- validate_transition(state, new_state, transitions, state_field) do
+      if function_exported?(fsm, :before_transition, 3) do
+        fsm.before_transition(struct, state, new_state)
+      else
+        fsm.before_transition(struct, state, new_state, state_field)
+      end
     end
   end
 
-  defp validate_transition(state, new_state, transitions) do
+  defp validate_transition(state, new_state, transitions, state_field) do
     transitions
     |> from_source_or_fallback(state)
     |> is_or_contains?(new_state)
     |> if do
       :ok
     else
-      {:error, "invalid transition from #{state} to #{new_state}"}
+      if state_field == Fsm.default_state_field() do
+        {:error, "invalid transition from #{state} to #{new_state}"}
+      else
+        {:error, "invalid transition from #{state} to #{new_state} for field #{state_field}"}
+      end
     end
   end
 
@@ -79,6 +105,9 @@ defmodule Fsmx do
 
   defp is_or_contains?(:*, _), do: true
   defp is_or_contains?(state, state), do: true
-  defp is_or_contains?(states, state) when is_list(states), do: Enum.member?(states, state) || Enum.member?(states, :*)
+
+  defp is_or_contains?(states, state) when is_list(states),
+    do: Enum.member?(states, state) || Enum.member?(states, :*)
+
   defp is_or_contains?(_, _), do: false
 end
diff --git a/lib/fsmx/fsm.ex b/lib/fsmx/fsm.ex
index ec8f847..a965df5 100644
--- a/lib/fsmx/fsm.ex
+++ b/lib/fsmx/fsm.ex
@@ -13,39 +13,59 @@ defmodule Fsmx.Fsm do
   defmodule MyApp.Fsm do
     use Fsmx.Fsm, transitions: %{}
 
-    def before_transition(struct, _from, _to) do
+    def before_transition(struct, _from, _to, _state_field) do
       # ...
     end
   end
+  ```
+
+  ## Callbacks
+
+  Callbacks are defined as functions in the module that includes `Fsmx.Fsm`
+  (which could be the one that includes `Fsmx.Struct`). They are called with
+  the struct, the current state, the new state, and the state field name. This
+  allows for different callbacks per state field.
+
+  The following callbacks are deprecated and will be removed in a future:
+  - `before_transition/3`
+  - `transition_changeset/4`
+  - `after_transition_multi/3`
   """
 
-  @callback before_transition(struct, Fsmx.state_t, Fsmx.state_t) :: {:ok, struct} | {:error, any}
+  @callback before_transition(struct, Fsmx.state_t(), Fsmx.state_t(), atom()) ::
+              {:ok, struct} | {:error, any}
 
   if Code.ensure_loaded?(Ecto) do
-    @callback transition_changeset(struct, Fsmx.state_t, Fsmx.state_t) :: Ecto.Changeset.t()
-    @callback after_transition_multi(struct, Fsmx.state_t, Fsmx.state_t) ::
+    @callback transition_changeset(Ecto.Schema.t(), Fsmx.state_t(), Fsmx.state_t(), map(), atom()) ::
+                Ecto.Changeset.t()
+    @callback after_transition_multi(struct, Fsmx.state_t(), Fsmx.state_t(), atom()) ::
                 {:ok, struct} | {:error, any}
   end
 
+  def default_state_field do
+    :state
+  end
+
   defmacro __using__(opts \\ []) do
+    {state_field, _} = Code.eval_quoted(Keyword.get(opts, :state_field, :state))
+
     quote do
       @before_compile unquote(__MODULE__)
 
-      @fsm Keyword.get(unquote(opts), :fsm, __MODULE__)
+      def __fsmx__(unquote(state_field), :transitions),
+        do: Keyword.fetch!(unquote(opts), :transitions)
 
-      def __fsmx__(:state_field), do: Keyword.get(unquote(opts), :state_field, :state)
-      def __fsmx__(:transitions), do: Keyword.fetch!(unquote(opts), :transitions)
-      def __fsmx__(:fsm), do: @fsm
+      def __fsmx__(unquote(state_field), :fsm), do: Keyword.get(unquote(opts), :fsm, __MODULE__)
     end
   end
 
   defmacro __before_compile__(_env) do
     quote generated: false do
-      def before_transition(struct, _from, _to), do: {:ok, struct}
+      def before_transition(struct, _from, _to, _state_field), do: {:ok, struct}
 
       if Code.ensure_loaded?(Ecto) do
-        def transition_changeset(changeset, _from, _to, _params), do: changeset
-        def after_transition_multi(struct, _from, _to), do: {:ok, struct}
+        def transition_changeset(changeset, _from, _to, _params, _state_field), do: changeset
+        def after_transition_multi(struct, _from, _to, _state_field), do: {:ok, struct}
       end
     end
   end
diff --git a/lib/fsmx/struct.ex b/lib/fsmx/struct.ex
index 760ed79..1869fa5 100644
--- a/lib/fsmx/struct.ex
+++ b/lib/fsmx/struct.ex
@@ -2,7 +2,7 @@ defmodule Fsmx.Struct do
   @moduledoc """
   Main module to include finite-state machine logic into your struct/schema
 
-  It assumes a `:state` string field exists in your model
+  If no `state_field` is defined, it assumes the name is `:state`.
 
   Basic usage:
 
@@ -13,17 +13,42 @@ defmodule Fsmx.Struct do
     use Fsmx.Struct, transitions: %{}
   end
   ```
+
+  You can also specify a custom state field:
+
+  ```elixir
+  defmodule MyApp.Struct do
+    defstruct [:my_state]
+
+    use Fsmx.Struct, state_field: :my_state, transitions: %{}
+  end
+  ```
+
+  Or even multiple state fields, that will behave independently and have their
+  own transition definition, etc. In this case `:state` is still used as the
+  default:
+
+  ```elixir
+  defmodule MyApp.Struct do
+    defstruct [:state, :other_state]
+
+    use Fsmx.Struct, transitions: %{}
+    use Fsmx.Struct, state_field: :other_state, transitions: %{}
+  end
+  ```
   """
 
   defmacro __using__(opts \\ []) do
+    {state_field, _} = Code.eval_quoted(Keyword.get(opts, :state_field, :state))
+
     quote do
-      @fsm Keyword.get(unquote(opts), :fsm, __MODULE__)
+      fsm = Keyword.get(unquote(opts), :fsm, __MODULE__)
 
-      if @fsm == __MODULE__ do
+      if fsm == __MODULE__ do
         use Fsmx.Fsm, unquote(opts)
       end
 
-      def __fsmx__, do: @fsm
+      def __fsmx__(unquote(state_field)), do: Keyword.get(unquote(opts), :fsm, __MODULE__)
     end
   end
 end
diff --git a/mix.exs b/mix.exs
index 52482e4..f7c613d 100644
--- a/mix.exs
+++ b/mix.exs
@@ -1,13 +1,13 @@
 defmodule Fsmx.MixProject do
   use Mix.Project
 
-  @version "0.4.1"
+  @version "0.5.0"
 
   def project do
     [
       app: :fsmx,
       version: @version,
-      elixir: "~> 1.8",
+      elixir: "~> 1.11",
       elixirc_paths: elixirc_paths(Mix.env()),
       start_permanent: Mix.env() == :prod,
       applications: applications(Mix.env()),
diff --git a/mix.lock b/mix.lock
index 0d4e6fe..a7c3a2f 100644
--- a/mix.lock
+++ b/mix.lock
@@ -1,11 +1,11 @@
 %{
   "connection": {:hex, :connection, "1.1.0", "ff2a49c4b75b6fb3e674bfc5536451607270aac754ffd1bdfe175abe4a6d7a68", [:mix], [], "hexpm", "722c1eb0a418fbe91ba7bd59a47e28008a189d47e37e0e7bb85585a016b2869c"},
   "db_connection": {:hex, :db_connection, "2.3.1", "4c9f3ed1ef37471cbdd2762d6655be11e38193904d9c5c1c9389f1b891a3088e", [:mix], [{:connection, "~> 1.0", [hex: :connection, repo: "hexpm", optional: false]}], "hexpm", "abaab61780dde30301d840417890bd9f74131041afd02174cf4e10635b3a63f5"},
-  "decimal": {:hex, :decimal, "2.0.0", "a78296e617b0f5dd4c6caf57c714431347912ffb1d0842e998e9792b5642d697", [:mix], [], "hexpm", "34666e9c55dea81013e77d9d87370fe6cb6291d1ef32f46a1600230b1d44f577"},
+  "decimal": {:hex, :decimal, "2.1.1", "5611dca5d4b2c3dd497dec8f68751f1f1a54755e8ed2a966c2633cf885973ad6", [:mix], [], "hexpm", "53cfe5f497ed0e7771ae1a475575603d77425099ba5faef9394932b35020ffcc"},
   "dialyxir": {:hex, :dialyxir, "1.0.0", "6a1fa629f7881a9f5aaf3a78f094b2a51a0357c843871b8bc98824e7342d00a5", [:mix], [{:erlex, ">= 0.2.6", [hex: :erlex, repo: "hexpm", optional: false]}], "hexpm", "aeb06588145fac14ca08d8061a142d52753dbc2cf7f0d00fc1013f53f8654654"},
   "earmark": {:hex, :earmark, "1.4.5", "62ffd3bd7722fb7a7b1ecd2419ea0b458c356e7168c1f5d65caf09b4fbdd13c8", [:mix], [], "hexpm", "b7d0e6263d83dc27141a523467799a685965bf8b13b6743413f19a7079843f4f"},
-  "ecto": {:hex, :ecto, "3.5.5", "48219a991bb86daba6e38a1e64f8cea540cded58950ff38fbc8163e062281a07", [:mix], [{:decimal, "~> 1.6 or ~> 2.0", [hex: :decimal, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: true]}, {:telemetry, "~> 0.4", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "98dd0e5e1de7f45beca6130d13116eae675db59adfa055fb79612406acf6f6f1"},
-  "ecto_sql": {:hex, :ecto_sql, "3.5.3", "1964df0305538364b97cc4661a2bd2b6c89d803e66e5655e4e55ff1571943efd", [:mix], [{:db_connection, "~> 2.2", [hex: :db_connection, repo: "hexpm", optional: false]}, {:ecto, "~> 3.5.0", [hex: :ecto, repo: "hexpm", optional: false]}, {:myxql, "~> 0.3.0 or ~> 0.4.0", [hex: :myxql, repo: "hexpm", optional: true]}, {:postgrex, "~> 0.15.0", [hex: :postgrex, repo: "hexpm", optional: true]}, {:tds, "~> 2.1.1", [hex: :tds, repo: "hexpm", optional: true]}, {:telemetry, "~> 0.4.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "d2f53592432ce17d3978feb8f43e8dc0705e288b0890caf06d449785f018061c"},
+  "ecto": {:hex, :ecto, "3.7.2", "44c034f88e1980754983cc4400585970b4206841f6f3780967a65a9150ef09a8", [:mix], [{:decimal, "~> 1.6 or ~> 2.0", [hex: :decimal, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: true]}, {:telemetry, "~> 0.4 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "a600da5772d1c31abbf06f3e4a1ffb150e74ed3e2aa92ff3cee95901657a874e"},
+  "ecto_sql": {:hex, :ecto_sql, "3.7.2", "55c60aa3a06168912abf145c6df38b0295c34118c3624cf7a6977cd6ce043081", [:mix], [{:db_connection, "~> 2.2", [hex: :db_connection, repo: "hexpm", optional: false]}, {:ecto, "~> 3.7.0", [hex: :ecto, repo: "hexpm", optional: false]}, {:myxql, "~> 0.4.0 or ~> 0.5.0 or ~> 0.6.0", [hex: :myxql, repo: "hexpm", optional: true]}, {:postgrex, "~> 0.15.0 or ~> 0.16.0 or ~> 1.0", [hex: :postgrex, repo: "hexpm", optional: true]}, {:tds, "~> 2.1.1 or ~> 2.2", [hex: :tds, repo: "hexpm", optional: true]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "3c218ea62f305dcaef0b915fb56583195e7b91c91dcfb006ba1f669bfacbff2a"},
   "erlex": {:hex, :erlex, "0.2.6", "c7987d15e899c7a2f34f5420d2a2ea0d659682c06ac607572df55a43753aa12e", [:mix], [], "hexpm", "2ed2e25711feb44d52b17d2780eabf998452f6efda104877a3881c2f8c0c0c75"},
   "etso": {:hex, :etso, "0.1.1", "bfc5e30483d397774a64981fc93511d3ed0dcb9d19bc3cba03df7b9555a68636", [:mix], [{:ecto, "~> 3.0", [hex: :ecto, repo: "hexpm", optional: false]}], "hexpm", "194ed738ce2e1197b326071b42dc092d6df4a7492575a63d6b797a0c59727088"},
   "ex_doc": {:hex, :ex_doc, "0.22.1", "9bb6d51508778193a4ea90fa16eac47f8b67934f33f8271d5e1edec2dc0eee4c", [:mix], [{:earmark, "~> 1.4.0", [hex: :earmark, repo: "hexpm", optional: false]}, {:makeup_elixir, "~> 0.14", [hex: :makeup_elixir, repo: "hexpm", optional: false]}], "hexpm", "d957de1b75cb9f78d3ee17820733dc4460114d8b1e11f7ee4fd6546e69b1db60"},
@@ -13,5 +13,5 @@
   "makeup_elixir": {:hex, :makeup_elixir, "0.14.1", "4f0e96847c63c17841d42c08107405a005a2680eb9c7ccadfd757bd31dabccfb", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "f2438b1a80eaec9ede832b5c41cd4f373b38fd7aa33e3b22d9db79e640cbde11"},
   "nimble_parsec": {:hex, :nimble_parsec, "0.6.0", "32111b3bf39137144abd7ba1cce0914533b2d16ef35e8abc5ec8be6122944263", [:mix], [], "hexpm", "27eac315a94909d4dc68bc07a4a83e06c8379237c5ea528a9acff4ca1c873c52"},
   "postgrex": {:hex, :postgrex, "0.15.7", "724410acd48abac529d0faa6c2a379fb8ae2088e31247687b16cacc0e0883372", [:mix], [{:connection, "~> 1.0", [hex: :connection, repo: "hexpm", optional: false]}, {:db_connection, "~> 2.1", [hex: :db_connection, repo: "hexpm", optional: false]}, {:decimal, "~> 1.5 or ~> 2.0", [hex: :decimal, repo: "hexpm", optional: false]}, {:jason, "~> 1.0", [hex: :jason, repo: "hexpm", optional: true]}], "hexpm", "88310c010ff047cecd73d5ceca1d99205e4b1ab1b9abfdab7e00f5c9d20ef8f9"},
-  "telemetry": {:hex, :telemetry, "0.4.2", "2808c992455e08d6177322f14d3bdb6b625fbcfd233a73505870d8738a2f4599", [:rebar3], [], "hexpm", "2d1419bd9dda6a206d7b5852179511722e2b18812310d304620c7bd92a13fcef"},
+  "telemetry": {:hex, :telemetry, "0.4.3", "a06428a514bdbc63293cd9a6263aad00ddeb66f608163bdec7c8995784080818", [:rebar3], [], "hexpm", "eb72b8365ffda5bed68a620d1da88525e326cb82a75ee61354fc24b844768041"},
 }
diff --git a/priv/repo/migrations/20200624134635_create_test_schemas.exs b/priv/repo/migrations/20200624134635_create_test_schemas.exs
index 3a30d3a..6c3cd41 100644
--- a/priv/repo/migrations/20200624134635_create_test_schemas.exs
+++ b/priv/repo/migrations/20200624134635_create_test_schemas.exs
@@ -5,6 +5,7 @@ defmodule Fsmx.Repo.Migrations.CreateTestSchemas do
     create table(:test) do
       add :state, :string
       add :before, :string
+      add :other_state, :string
     end
   end
 end
diff --git a/test/fsmx/ecto_test.exs b/test/fsmx/ecto_test.exs
index a314917..e997e59 100644
--- a/test/fsmx/ecto_test.exs
+++ b/test/fsmx/ecto_test.exs
@@ -3,7 +3,7 @@ defmodule Fsmx.EctoTest do
 
   alias Ecto.Multi
   alias Fsmx.Repo
-  alias Fsmx.TestEctoSchemas.{Simple, WithCallbacks, WithSeparateFsm}
+  alias Fsmx.TestEctoSchemas.{Simple, WithCallbacks, WithSeparateFsm, MultiState}
 
   describe "transition_changeset/2" do
     test "returns a changeset" do
@@ -37,10 +37,18 @@ defmodule Fsmx.EctoTest do
 
       assert Ecto.Changeset.get_change(two_changeset, :state) == "2"
     end
+
+    test "works the same with multiple states" do
+      one = %MultiState{state: "1", other_state: "1"}
+
+      two_changeset = Fsmx.transition_changeset(one, "2", [], state_field: :other_state)
+
+      assert Ecto.Changeset.get_change(two_changeset, :other_state) == "2"
+    end
   end
 
   describe "transition/2 with callbacks" do
-    test "calls before_transition/2 on struct" do
+    test "calls before_transition/3 on struct" do
       one = %WithCallbacks.ValidBefore{state: "1", before: false}
 
       two = Fsmx.transition_changeset(one, "2")
@@ -48,7 +56,7 @@ defmodule Fsmx.EctoTest do
       assert %WithCallbacks.ValidBefore{before: "1"} = two.data
     end
 
-    test "fails if before_transition/2 returns an error" do
+    test "fails if before_transition/3 returns an error" do
       one = %WithCallbacks.InvalidBefore{state: "1", before: false}
 
       changeset = Fsmx.transition_changeset(one, "2")
@@ -56,6 +64,16 @@ defmodule Fsmx.EctoTest do
       refute changeset.valid?
       assert changeset.errors == [state: {"transition_changeset failed: before_failed", []}]
     end
+
+    test "call before_transition/4 on struct with new state" do
+      one = %WithCallbacks.MultiStateValidBefore{state: "1", other_state: "1", before: false}
+
+      two = Fsmx.transition_changeset(one, "2")
+      assert %WithCallbacks.MultiStateValidBefore{before: "1"} = two.data
+
+      new_two = Fsmx.transition_changeset(one, "2", %{}, state_field: :other_state)
+      assert %WithCallbacks.MultiStateValidBefore{before: "2"} = new_two.data
+    end
   end
 
   describe "transition/2 with separate fsm module" do
@@ -132,5 +150,19 @@ defmodule Fsmx.EctoTest do
 
       assert %WithCallbacks.InvalidAfterMulti{state: "1"} = updated_schema
     end
+
+    test "adds a transition changeset to the given multi for a new state" do
+      one = %MultiState{state: "1", other_state: "1"}
+
+      multi =
+        Fsmx.transition_multi(Multi.new(), one, "transition", "2", %{}, state_field: :other_state)
+
+      assert %Ecto.Multi{operations: operations} = multi
+
+      assert [_, {"transition", {:changeset, two_changeset, []}}] = operations
+      assert %Ecto.Changeset{} = two_changeset
+      assert two_changeset.data.other_state == "1"
+      assert Ecto.Changeset.get_change(two_changeset, :other_state) == "2"
+    end
   end
 end
diff --git a/test/fsmx/struct_test.exs b/test/fsmx/struct_test.exs
index cea4ce9..9122cfb 100644
--- a/test/fsmx/struct_test.exs
+++ b/test/fsmx/struct_test.exs
@@ -1,7 +1,7 @@
 defmodule Fsmx.StructTest do
   use ExUnit.Case
 
-  alias Fsmx.TestStructs.{Simple, WithCallbacks, WithSeparateFsm, WithFallback}
+  alias Fsmx.TestStructs.{Simple, WithCallbacks, WithSeparateFsm, WithFallback, MultiState}
 
   describe "transition/2" do
     test "can do simple transitions" do
@@ -67,4 +67,40 @@ defmodule Fsmx.StructTest do
       assert %WithSeparateFsm{state: "2", before: "1"} = two
     end
   end
+
+  describe "transition/2 with multiple state fields" do
+    test "the default state works the same" do
+      one = %MultiState{state: "1", other_state: "1"}
+
+      {:ok, two} = Fsmx.transition(one, "2")
+
+      assert %MultiState{state: "2", other_state: "1"} = two
+    end
+
+    test "transitioning the new state" do
+      one = %MultiState{state: "1", other_state: "1"}
+
+      {:ok, two} = Fsmx.transition(one, "2", state_field: :other_state)
+
+      assert %MultiState{state: "1", other_state: "2"} = two
+    end
+
+    test "fails to perform invalid transitions on the new state" do
+      one = %MultiState{state: "1", other_state: "1"}
+
+      assert {:error, msg} = Fsmx.transition(one, "3", state_field: :other_state)
+
+      assert msg == "invalid transition from 1 to 3 for field other_state"
+    end
+
+    test ":* as destination means the state can transit to any other state using the new state" do
+      three = %MultiState{state: "3", other_state: "3"}
+
+      assert {:ok, %{state: "3", other_state: "1"}} =
+               Fsmx.transition(three, "1", state_field: :other_state)
+
+      assert {:ok, %{state: "3", other_state: "2"}} =
+               Fsmx.transition(three, "2", state_field: :other_state)
+    end
+  end
 end
diff --git a/test/support/test_ecto_schemas/multi_state.ex b/test/support/test_ecto_schemas/multi_state.ex
new file mode 100644
index 0000000..ebdafab
--- /dev/null
+++ b/test/support/test_ecto_schemas/multi_state.ex
@@ -0,0 +1,12 @@
+defmodule Fsmx.TestEctoSchemas.MultiState do
+  use Ecto.Schema
+
+  schema "test" do
+    field :state, :string
+    field :other_state, :string
+  end
+
+  use Fsmx.Struct, transitions: %{"1" => "2"}
+
+  use Fsmx.Struct, state_field: :other_state, transitions: %{"1" => "2"}
+end
diff --git a/test/support/test_ecto_schemas/with_callbacks.ex b/test/support/test_ecto_schemas/with_callbacks.ex
index 27119d7..3efcde2 100644
--- a/test/support/test_ecto_schemas/with_callbacks.ex
+++ b/test/support/test_ecto_schemas/with_callbacks.ex
@@ -57,4 +57,25 @@ defmodule Fsmx.TestEctoSchemas.WithCallbacks do
       {:error, :after_transition_multi_failed}
     end
   end
+
+  defmodule MultiStateValidBefore do
+    use Ecto.Schema
+
+    schema "test" do
+      field :state, :string
+      field :other_state, :string
+      field :before, :string
+    end
+
+    use Fsmx.Struct, transitions: %{"1" => "2"}
+    use Fsmx.Struct, state_field: :other_state, transitions: %{"1" => "2"}
+
+    def before_transition(struct, "1", _new_state, :state) do
+      {:ok, %{struct | before: "1"}}
+    end
+
+    def before_transition(struct, "1", _new_state, :other_state) do
+      {:ok, %{struct | before: "2"}}
+    end
+  end
 end
diff --git a/test/support/test_structs/multi_state.ex b/test/support/test_structs/multi_state.ex
new file mode 100644
index 0000000..49f3994
--- /dev/null
+++ b/test/support/test_structs/multi_state.ex
@@ -0,0 +1,10 @@
+defmodule Fsmx.TestStructs.MultiState do
+  defstruct [:state, :other_state]
+
+  use Fsmx.Struct,
+    transitions: %{"1" => ["2"], "2" => ["3"], "3" => :*}
+
+  use Fsmx.Struct,
+    state_field: :other_state,
+    transitions: %{"1" => ["2"], "2" => ["3"], "3" => :*}
+end