diff --git a/lib/fsmx/fsm.ex b/lib/fsmx/fsm.ex
index f47a64a..30a5a2b 100644
--- a/lib/fsmx/fsm.ex
+++ b/lib/fsmx/fsm.ex
@@ -36,6 +36,8 @@ defmodule Fsmx.Fsm do
               {:ok, struct} | {:error, any}
 
   if Code.ensure_loaded?(Ecto) do
+    @callback transition_changeset(Ecto.Schema.t(), Fsmx.state_t(), Fsmx.state_t(), map()) ::
+                Ecto.Changeset.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()) ::
@@ -50,6 +52,11 @@ defmodule Fsmx.Fsm do
     {state_field, _} = Code.eval_quoted(Keyword.get(opts, :state_field, :state))
 
     quote do
+      if not Module.get_attribute(__MODULE__, :__fsmx__behaviour, false) do
+        @behaviour Fsmx.Fsm
+        Module.put_attribute(__MODULE__, :__fsmx__behaviour, true)
+      end
+
       @before_compile unquote(__MODULE__)
 
       def __fsmx__(unquote(state_field), :transitions),
@@ -61,11 +68,17 @@ defmodule Fsmx.Fsm do
 
   defmacro __before_compile__(_env) do
     quote generated: false do
+      @impl Fsmx.Fsm
       def before_transition(struct, _from, _to, _state_field), do: {:ok, struct}
 
       if Code.ensure_loaded?(Ecto) do
+        @impl Fsmx.Fsm
         def transition_changeset(changeset, _from, _to, _params), do: changeset
+
+        @impl Fsmx.Fsm
         def transition_changeset(changeset, _from, _to, _params, _state_field), do: changeset
+
+        @impl Fsmx.Fsm
         def after_transition_multi(struct, _from, _to, _state_field), do: {:ok, struct}
       end
     end