2 files changed, 10 insertions, 4 deletions
diff --git a/test/web/mastodon_api/views/notification_view_test.exs b/test/web/mastodon_api/views/notification_view_test.exs
index 384fe7253..c5691341a 100644
--- a/test/web/mastodon_api/views/notification_view_test.exs
+++ b/test/web/mastodon_api/views/notification_view_test.exs
@@ -7,6 +7,7 @@ defmodule Pleroma.Web.MastodonAPI.NotificationViewTest do
 
   alias Pleroma.Activity
   alias Pleroma.Chat
+  alias Pleroma.ChatMessageReference
   alias Pleroma.Notification
   alias Pleroma.Object
   alias Pleroma.Repo
@@ -16,7 +17,7 @@ defmodule Pleroma.Web.MastodonAPI.NotificationViewTest do
   alias Pleroma.Web.MastodonAPI.AccountView
   alias Pleroma.Web.MastodonAPI.NotificationView
   alias Pleroma.Web.MastodonAPI.StatusView
-  alias Pleroma.Web.PleromaAPI.ChatMessageView
+  alias Pleroma.Web.PleromaAPI.ChatMessageReferenceView
   import Pleroma.Factory
 
   defp test_notifications_rendering(notifications, user, expected_result) do
@@ -44,13 +45,15 @@ defmodule Pleroma.Web.MastodonAPI.NotificationViewTest do
     object = Object.normalize(activity)
     chat = Chat.get(recipient.id, user.ap_id)
 
+    cm_ref = ChatMessageReference.for_chat_and_object(chat, object)
+
     expected = %{
       id: to_string(notification.id),
       pleroma: %{is_seen: false},
       type: "pleroma:chat_mention",
       account: AccountView.render("show.json", %{user: user, for: recipient}),
       chat_message:
-        ChatMessageView.render("show.json", %{object: object, for: recipient, chat: chat}),
+        ChatMessageReferenceView.render("show.json", %{chat_message_reference: cm_ref}),
       created_at: Utils.to_masto_date(notification.inserted_at)
     }
 
diff --git a/test/web/pleroma_api/views/chat_view_test.exs b/test/web/pleroma_api/views/chat_view_test.exs
index 6062a0cfe..f3bd12616 100644
--- a/test/web/pleroma_api/views/chat_view_test.exs
+++ b/test/web/pleroma_api/views/chat_view_test.exs
@@ -6,11 +6,12 @@ defmodule Pleroma.Web.PleromaAPI.ChatViewTest do
   use Pleroma.DataCase
 
   alias Pleroma.Chat
+  alias Pleroma.ChatMessageReference
   alias Pleroma.Object
   alias Pleroma.Web.CommonAPI
   alias Pleroma.Web.CommonAPI.Utils
   alias Pleroma.Web.MastodonAPI.AccountView
-  alias Pleroma.Web.PleromaAPI.ChatMessageView
+  alias Pleroma.Web.PleromaAPI.ChatMessageReferenceView
   alias Pleroma.Web.PleromaAPI.ChatView
 
   import Pleroma.Factory
@@ -39,7 +40,9 @@ defmodule Pleroma.Web.PleromaAPI.ChatViewTest do
 
     represented_chat = ChatView.render("show.json", chat: chat)
 
+    cm_ref = ChatMessageReference.for_chat_and_object(chat, chat_message)
+
     assert represented_chat[:last_message] ==
-             ChatMessageView.render("show.json", chat: chat, object: chat_message)
+             ChatMessageReferenceView.render("show.json", chat_message_reference: cm_ref)
   end
 end