From 5fe0325e35fb502f5180eebe3115ad478b9da627 Mon Sep 17 00:00:00 2001
From: Eileen Noonan <enoonan@arcstone.com>
Date: Mon, 25 Nov 2024 15:39:33 -0500
Subject: [PATCH] enforce valid upsert transition targets at runtime

---
 lib/ash_state_machine.ex        | 56 +++++++++++++++++++++++----------
 test/ash_state_machine_test.exs |  6 ++++
 test/support/verification.ex    |  7 +++++
 3 files changed, 52 insertions(+), 17 deletions(-)

diff --git a/lib/ash_state_machine.ex b/lib/ash_state_machine.ex
index 49472d4..ff05273 100644
--- a/lib/ash_state_machine.ex
+++ b/lib/ash_state_machine.ex
@@ -137,20 +137,27 @@ defmodule AshStateMachine do
     end
   end
 
+  def transition_state(%{action_type: :create, action: %{upsert?: true}} = changeset, target) do
+    attribute = AshStateMachine.Info.state_machine_state_attribute!(changeset.resource)
+    old_state = Map.get(changeset.data, attribute)
+
+    cond do
+      target not in AshStateMachine.Info.state_machine_initial_states!(changeset.resource) ->
+        invalid_initial_state(changeset, target)
+
+      target not in available_upsert_targets(changeset) -> no_matching_transition(changeset, target, old_state)
+
+      true -> Ash.Changeset.force_change_attribute(changeset, attribute, target)
+    end
+  end
+
   def transition_state(%{action_type: :create} = changeset, target) do
     attribute = AshStateMachine.Info.state_machine_state_attribute!(changeset.resource)
 
     if target in AshStateMachine.Info.state_machine_initial_states!(changeset.resource) do
       Ash.Changeset.force_change_attribute(changeset, attribute, target)
     else
-      changeset
-      |> Ash.Changeset.set_context(%{state_machine: %{attempted_change: target}})
-      |> Ash.Changeset.add_error(
-        AshStateMachine.Errors.InvalidInitialState.exception(
-          target: target,
-          action: changeset.action.name
-        )
-      )
+      invalid_initial_state(changeset, target)
     end
   end
 
@@ -158,6 +165,17 @@ defmodule AshStateMachine do
     Ash.Changeset.add_error(other, "Can't transition states on destroy actions")
   end
 
+  defp invalid_initial_state(changeset, target) do
+    changeset
+    |> Ash.Changeset.set_context(%{state_machine: %{attempted_change: target}})
+    |> Ash.Changeset.add_error(
+      AshStateMachine.Errors.InvalidInitialState.exception(
+        target: target,
+        action: changeset.action.name
+      )
+    )
+  end
+
   defp find_and_perform_transition(changeset, old_state, attribute, target) do
     changeset.resource
     |> AshStateMachine.Info.state_machine_transitions(changeset.action.name)
@@ -166,15 +184,7 @@ defmodule AshStateMachine do
     end)
     |> case do
       nil ->
-        changeset
-        |> Ash.Changeset.set_context(%{state_machine: %{attempted_change: target}})
-        |> Ash.Changeset.add_error(
-          AshStateMachine.Errors.NoMatchingTransition.exception(
-            old_state: old_state,
-            target: target,
-            action: changeset.action.name
-          )
-        )
+         no_matching_transition(changeset, target, old_state)
 
       _transition ->
         Ash.Changeset.force_change_attribute(changeset, attribute, target)
@@ -197,6 +207,11 @@ defmodule AshStateMachine do
       To remediate this, add the `extra_states` option and include the state #{inspect(target)}
     """)
 
+     no_matching_transition(changeset, target, old_state)
+  end
+
+  @doc false
+  defp no_matching_transition(changeset, target, old_state) do
     changeset
     |> Ash.Changeset.set_context(%{state_machine: %{attempted_change: target}})
     |> Ash.Changeset.add_error(
@@ -243,4 +258,11 @@ defmodule AshStateMachine do
     |> Enum.reject(&(&1 == :*))
     |> Enum.uniq()
   end
+
+  defp available_upsert_targets(changeset) do
+    AshStateMachine.Info.state_machine_transitions(changeset.resource, changeset.action.name)
+    |> Enum.map(& &1.to)
+    |> List.flatten()
+    |> Enum.uniq()
+  end
 end
diff --git a/test/ash_state_machine_test.exs b/test/ash_state_machine_test.exs
index b0ff7f8..a5bb553 100644
--- a/test/ash_state_machine_test.exs
+++ b/test/ash_state_machine_test.exs
@@ -50,6 +50,12 @@ defmodule AshStateMachineTest do
       assert Verification.reset!(%{id: state_machine.id}).state == :pending
     end
 
+    test "create upsert? actions do not allow invalid states" do
+      state_machine = Verification.create!() |> Verification.begin!()
+      assert {:error, reason} = Verification.broken_upsert(%{id: state_machine.id})
+      assert Exception.message(reason) =~ ~r/no matching transition/i
+    end
+
     test "create actions without `upsert? true` do not compile" do
       assert_raise Spark.Error.DslError, ~r/non-upsert create action/, fn ->
         defmodule CreateWithoutUpsert do
diff --git a/test/support/verification.ex b/test/support/verification.ex
index c1ac234..b7036a5 100644
--- a/test/support/verification.ex
+++ b/test/support/verification.ex
@@ -11,6 +11,7 @@ defmodule Verification do
     transitions do
       transition(:begin, from: :pending, to: :executing)
       transition(:reset, from: :*, to: :pending)
+      transition(:broken_upsert, from: :*, to: [:foo, :bar])
     end
   end
 
@@ -26,6 +27,11 @@ defmodule Verification do
       upsert? true
       change transition_state(:pending)
     end
+
+    create :broken_upsert do
+      upsert? true
+      change transition_state(:pending)
+    end
   end
 
   attributes do
@@ -36,5 +42,6 @@ defmodule Verification do
     define :create
     define :begin
     define :reset
+    define :broken_upsert
   end
 end