From 13de02c097a967e494f5ec21555646e1eabc403e Mon Sep 17 00:00:00 2001
From: Pyro <178290668+pyros-projects@users.noreply.github.com>
Date: Tue, 4 Aug 2026 01:41:02 +0200
Subject: [PATCH] Test MiniMax output attention patches

---
 tests-unit/comfy_test/test_minimax_attention.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/tests-unit/comfy_test/test_minimax_attention.py b/tests-unit/comfy_test/test_minimax_attention.py
index 3f010be24..97e9a6116 100644
--- a/tests-unit/comfy_test/test_minimax_attention.py
+++ b/tests-unit/comfy_test/test_minimax_attention.py
@@ -26,6 +26,10 @@ def test_attention_patch_accepts_tuple_and_mapping_callbacks(monkeypatch):
         assert extra_options["n_heads"] == 2
         return {"q": q * 2, "v": v * 3}
 
+    def output_patch(out, extra_options):
+        assert extra_options["block_index"] == 3
+        return out + 4
+
     def fake_attention(q, k, v, *args, **kwargs):
         seen.update(q=q, k=k, v=v)
         return v.transpose(1, 2).reshape(1, v.shape[2], -1)
@@ -36,11 +40,14 @@ def test_attention_patch_accepts_tuple_and_mapping_callbacks(monkeypatch):
         x,
         transformer_options={
             "block_index": 3,
-            "patches": {"attn1_patch": [tuple_patch, mapping_patch]},
+            "patches": {
+                "attn1_patch": [tuple_patch, mapping_patch],
+                "attn1_output_patch": [output_patch],
+            },
         },
     )
 
     assert torch.equal(seen["q"], torch.full((1, 2, 2, 2), 2.0))
     assert torch.equal(seen["k"], torch.full((1, 2, 2, 2), 2.0))
     assert torch.equal(seen["v"], torch.full((1, 2, 2, 2), 9.0))
-    assert torch.equal(output, torch.full((2, 4), 9.0))
+    assert torch.equal(output, torch.full((2, 4), 13.0))