Merge pull request rails#54371 from byroot/fix-on-rotate-callback

Fix `rotate(on_on_rotation:)` and `#on_rotation`

activesupport/lib/active_support/messages/rotator.rb

@@ -11,10 +11,16 @@ def initialize(*args, on_rotation: nil, **options)
         @on_rotation = on_rotation
       end
 
-      def rotate(*args, **options)
+      def rotate(*args, on_rotation: nil, **options)
+        @on_rotation = on_rotation if on_rotation
         fall_back_to build_rotation(*args, **options)
       end
 
+      def on_rotation(&on_rotation)
+        @on_rotation = on_rotation
+        self
+      end
+
       def fall_back_to(fallback)
         @rotations << fallback
         self

activesupport/test/messages/message_rotator_tests.rb

@@ -45,6 +45,26 @@ def self.load(*); raise Class.new(StandardError); end
       assert called
     end
 
+    test "rotate(on_rotation:) is called on successful rotation" do
+      called = nil
+      codec = make_codec(secret("new")).rotate(secret("old"), on_rotation: proc { called = true })
+      old_codec = make_codec(secret("old"))
+      old_message = encode(DATA, old_codec)
+      assert_equal DATA, decode(old_message, codec)
+      assert called
+    end
+
+    test "rotate().on_rotation is called on successful rotation" do
+      called = nil
+      codec = make_codec(secret("new")).rotate(secret("old")).on_rotation do
+        called = true
+      end
+      old_codec = make_codec(secret("old"))
+      old_message = encode(DATA, old_codec)
+      assert_equal DATA, decode(old_message, codec)
+      assert called
+    end
+
     test "on_rotation is not called when no rotation is necessary" do
       called = nil
       assert_rotate [secret("same"), on_rotation: proc { called = true }], [secret("same")]