2 files changed, 56 insertions, 2 deletions
diff --git a/test/pleroma/web/activity_pub/mrf/simple_policy_test.exs b/test/pleroma/web/activity_pub/mrf/simple_policy_test.exs
index 674b506c3..57fc00af5 100644
--- a/test/pleroma/web/activity_pub/mrf/simple_policy_test.exs
+++ b/test/pleroma/web/activity_pub/mrf/simple_policy_test.exs
@@ -57,6 +57,16 @@ defmodule Pleroma.Web.ActivityPub.MRF.SimplePolicyTest do
 
       assert SimplePolicy.filter(local_message) == {:ok, local_message}
     end
+
+    test "works with Updates" do
+      clear_config([:mrf_simple, :media_removal], [{"remote.instance", "Some reason"}])
+      media_message = build_media_message(type: "Update")
+
+      assert SimplePolicy.filter(media_message) ==
+               {:ok,
+                media_message
+                |> Map.put("object", Map.delete(media_message["object"], "attachment"))}
+    end
   end
 
   describe "when :media_nsfw" do
@@ -90,12 +100,20 @@ defmodule Pleroma.Web.ActivityPub.MRF.SimplePolicyTest do
 
       assert SimplePolicy.filter(local_message) == {:ok, local_message}
     end
+
+    test "works with Updates" do
+      clear_config([:mrf_simple, :media_nsfw], [{"remote.instance", "Whetever"}])
+      media_message = build_media_message(type: "Update")
+
+      assert SimplePolicy.filter(media_message) ==
+               {:ok, put_in(media_message, ["object", "sensitive"], true)}
+    end
   end
 
-  defp build_media_message do
+  defp build_media_message(opts \\ []) do
     %{
       "actor" => "https://remote.instance/users/bob",
-      "type" => "Create",
+      "type" => opts[:type] || "Create",
       "object" => %{
         "attachment" => [%{}],
         "tag" => ["foo"],
diff --git a/test/pleroma/web/activity_pub/mrf/tag_policy_test.exs b/test/pleroma/web/activity_pub/mrf/tag_policy_test.exs
index 46be316ee..a0db8df54 100644
--- a/test/pleroma/web/activity_pub/mrf/tag_policy_test.exs
+++ b/test/pleroma/web/activity_pub/mrf/tag_policy_test.exs
@@ -99,6 +99,24 @@ defmodule Pleroma.Web.ActivityPub.MRF.TagPolicyTest do
 
       assert TagPolicy.filter(message) == {:ok, except_message}
     end
+
+    test "removes attachments in Updates" do
+      actor = insert(:user, tags: ["mrf_tag:media-strip"])
+
+      message = %{
+        "actor" => actor.ap_id,
+        "type" => "Update",
+        "object" => %{"attachment" => ["file1"]}
+      }
+
+      except_message = %{
+        "actor" => actor.ap_id,
+        "type" => "Update",
+        "object" => %{}
+      }
+
+      assert TagPolicy.filter(message) == {:ok, except_message}
+    end
   end
 
   describe "mrf_tag:media-force-nsfw" do
@@ -119,5 +137,23 @@ defmodule Pleroma.Web.ActivityPub.MRF.TagPolicyTest do
 
       assert TagPolicy.filter(message) == {:ok, except_message}
     end
+
+    test "Mark as sensitive on presence of attachments in Updates" do
+      actor = insert(:user, tags: ["mrf_tag:media-force-nsfw"])
+
+      message = %{
+        "actor" => actor.ap_id,
+        "type" => "Update",
+        "object" => %{"tag" => ["test"], "attachment" => ["file1"]}
+      }
+
+      except_message = %{
+        "actor" => actor.ap_id,
+        "type" => "Update",
+        "object" => %{"tag" => ["test"], "attachment" => ["file1"], "sensitive" => true}
+      }
+
+      assert TagPolicy.filter(message) == {:ok, except_message}
+    end
   end
 end