From 23b75d8a2bb391da6dd803455576de9bdfed3c6c Mon Sep 17 00:00:00 2001
From: Roman Rizzi <rizziromanalejandro@gmail.com>
Date: Thu, 30 Dec 2021 15:40:16 -0300
Subject: [PATCH] FEATURE: Before consolidation callbacks. (#15428)

You can add callbacks that get called before updating an already consolidated notification or creating a consolidated one.

Instances of this rule can add callbacks to access the old notifications about to be destroyed or the consolidated one and add additional data inside the data hash versus having to execute extra queries when adding this logic inside the `set_mutations` block.
---
 .../consolidate_notifications.rb              | 23 +++++++
 .../consolidate_notifications_spec.rb         | 60 +++++++++++++++++++
 2 files changed, 83 insertions(+)
 create mode 100644 spec/services/notifications/consolidate_notifications_spec.rb

diff --git a/app/services/notifications/consolidate_notifications.rb b/app/services/notifications/consolidate_notifications.rb
index f9f1faf2281..adfeaee10fc 100644
--- a/app/services/notifications/consolidate_notifications.rb
+++ b/app/services/notifications/consolidate_notifications.rb
@@ -22,6 +22,15 @@
 # Need to call #set_mutations to configure this:
 #
 # - set_data_blk: A block that receives the notification data hash and mutates it, adding additional data needed for consolidation.
+#
+# Need to call #before_consolidation_callbacks to configure this:
+#
+# - before_update_blk: A block that is called before updating an already consolidated notification.
+#                      Receives the consolidated object, the data hash, and the original notification.
+#
+# - before_consolidation_blk: A block that is called before creating a consolidated object.
+#                             Receives an ActiveRecord::Relation with notifications about to be consolidated, and the new data hash.
+#
 
 module Notifications
   class ConsolidateNotifications < ConsolidationPlan
@@ -37,6 +46,12 @@ module Notifications
       @bump_notification = bump_notification
     end
 
+    def before_consolidation_callbacks(before_update_blk: nil, before_consolidation_blk: nil)
+      @before_update_blk = before_update_blk
+      @before_consolidation_blk = before_consolidation_blk
+      self
+    end
+
     def can_consolidate_data?(notification)
       return false if get_threshold.zero? || to.blank?
       return false if notification.notification_type != from
@@ -75,6 +90,10 @@ module Notifications
       data_hash = consolidated.data_hash.merge(data)
       data_hash[:count] += 1 if data_hash[:count].present?
 
+      if @before_update_blk
+        @before_update_blk.call(consolidated, data_hash, notification)
+      end
+
       # Hack: We don't want to cache the old data if we're about to update it.
       consolidated.instance_variable_set(:@data_hash, nil)
 
@@ -100,6 +119,10 @@ module Notifications
       timestamp = notifications.last.created_at
       data[:count] = count_after_saving_notification
 
+      if @before_consolidation_blk
+        @before_consolidation_blk.call(notifications, data)
+      end
+
       consolidated = nil
 
       Notification.transaction do
diff --git a/spec/services/notifications/consolidate_notifications_spec.rb b/spec/services/notifications/consolidate_notifications_spec.rb
new file mode 100644
index 00000000000..5ac150246a9
--- /dev/null
+++ b/spec/services/notifications/consolidate_notifications_spec.rb
@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+require 'rails_helper'
+
+describe Notifications::ConsolidateNotifications do
+  describe '#before_consolidation_callbacks' do
+    fab!(:user) { Fabricate(:user) }
+    let(:rule) do
+      described_class.new(
+        from: Notification.types[:liked],
+        to: Notification.types[:liked],
+        consolidation_window: 10.minutes,
+        consolidated_query_blk: Proc.new do |notifications|
+          notifications.where("(data::json ->> 'consolidated')::bool")
+        end,
+        threshold: 1
+      ).set_mutations(set_data_blk: Proc.new { |n| n.data_hash.merge(consolidated: true) })
+    end
+
+    it 'applies a callback when consolidating a notification' do
+      rule.before_consolidation_callbacks(
+        before_consolidation_blk: Proc.new do |_, data|
+          data[:consolidation_callback_called] = true
+        end
+      )
+
+      rule.consolidate_or_save!(build_like_notification)
+      rule.consolidate_or_save!(build_like_notification)
+
+      consolidated_notification = Notification.where(user: user).last
+
+      expect(consolidated_notification.data_hash[:consolidation_callback_called]).to eq(true)
+    end
+
+    it 'applies a callback when updating a consolidated notification' do
+      rule.before_consolidation_callbacks(
+        before_update_blk: Proc.new do |_, data|
+          data[:update_callback_called] = true
+        end
+      )
+
+      rule.consolidate_or_save!(build_like_notification)
+      rule.consolidate_or_save!(build_like_notification)
+
+      consolidated_notification = Notification.where(user: user).last
+
+      expect(consolidated_notification.data_hash[:update_callback_called]).to be_nil
+
+      rule.consolidate_or_save!(build_like_notification)
+
+      consolidated_notification = Notification.where(user: user).last
+
+      expect(consolidated_notification.data_hash[:update_callback_called]).to eq(true)
+    end
+
+    def build_like_notification
+      Fabricate.build(:notification, user: user, notification_type: Notification.types[:liked], data: {}.to_json)
+    end
+  end
+end