From a7d069ae3086ae6e5badd0ee51e598cac115a772 Mon Sep 17 00:00:00 2001
From: Zach Daniel <zach@zachdaniel.dev>
Date: Wed, 14 Jun 2023 20:40:04 -0400
Subject: [PATCH] fix: properly set notification_data from loaded record

---
 lib/ash/actions/helpers.ex      | 10 ++++++++--
 lib/ash/actions/load.ex         |  4 ++--
 test/notifier/notifier_test.exs | 10 ++++++++++
 3 files changed, 20 insertions(+), 4 deletions(-)

diff --git a/lib/ash/actions/helpers.ex b/lib/ash/actions/helpers.ex
index cec901d5..a23ea924 100644
--- a/lib/ash/actions/helpers.ex
+++ b/lib/ash/actions/helpers.ex
@@ -435,7 +435,13 @@ defmodule Ash.Actions.Helpers do
 
       case api.load(result, query, opts) do
         {:ok, result} ->
-          {:ok, result, Map.put(instructions, :notification_data, result)}
+          {:ok, result,
+           Map.update(
+             instructions,
+             :set_keys,
+             %{notification_data: result},
+             &Map.put(&1, :notification_data, result)
+           )}
 
         {:error, error} ->
           {:error, error}
@@ -454,7 +460,7 @@ defmodule Ash.Actions.Helpers do
 
       case api.load(result, query, opts) do
         {:ok, result} ->
-          {:ok, result}
+          {:ok, result, %{set_keys: %{notification_data: result}}}
 
         {:error, error} ->
           {:error, error}
diff --git a/lib/ash/actions/load.ex b/lib/ash/actions/load.ex
index 67a24c10..08926f0a 100644
--- a/lib/ash/actions/load.ex
+++ b/lib/ash/actions/load.ex
@@ -690,8 +690,8 @@ defmodule Ash.Actions.Load do
          parent_data_path,
          _parent_query_path,
          _join_request_path,
-         actual_query,
-         _
+         _,
+         actual_query
        )
        when not is_nil(manual) do
     {mod, opts} =
diff --git a/test/notifier/notifier_test.exs b/test/notifier/notifier_test.exs
index dfbd58a7..8cce0240 100644
--- a/test/notifier/notifier_test.exs
+++ b/test/notifier/notifier_test.exs
@@ -80,6 +80,8 @@ defmodule Ash.Test.NotifierTest do
       defaults [:create, :read, :update, :destroy]
 
       create :create_with_comment do
+        change load(:comments)
+
         change fn changeset, _ ->
           Ash.Changeset.after_action(changeset, fn _changeset, result ->
             Comment
@@ -191,6 +193,14 @@ defmodule Ash.Test.NotifierTest do
     assert_receive {:notification, %Ash.Notifier.Notification{data: %Comment{name: "auto"}}}
   end
 
+  test "the `load/1` change puts the loaded data into the notification" do
+    Post
+    |> Ash.Changeset.for_create(:create_with_comment, %{name: "foobar"})
+    |> Api.create!()
+
+    assert_receive {:notification, %Ash.Notifier.Notification{data: %Post{comments: [_]}}}
+  end
+
   test "notifications use the data before its limited by a select statement" do
     Comment
     |> Ash.Changeset.new(%{name: "foobar"})