optimizely · msohailhussain · Oct 25, 2019 · Oct 25, 2019 · Oct 25, 2019 · Oct 25, 2019
diff --git a/lib/optimizely/notification_center.rb b/lib/optimizely/notification_center.rb
@@ -39,20 +39,28 @@ def initialize(logger, error_handler)
 
     # Adds notification callback to the notification center
     #
-    # @param notification_type -  One of the constants in NOTIFICATION_TYPES
-    # @param notification_callback -  Function to call when the event is sent
+    # @param notification_type - One of the constants in NOTIFICATION_TYPES
+    # @param notification_callback [lambda, Method, Callable] (default: block) - Called when the event is sent
+    # @yield Block to be used as callback if callback omitted.
     #
     # @return [notification ID] Used to remove the notification
 
-    def add_notification_listener(notification_type, notification_callback)
+    def add_notification_listener(notification_type, notification_callback = nil, &block)
       return nil unless notification_type_valid?(notification_type)
 
+      if notification_callback && block_given?
+        @logger.log Logger::ERROR, 'Callback and block are mutually exclusive.'
+        return nil
+      end
+
+      notification_callback ||= block
+
       unless notification_callback
         @logger.log Logger::ERROR, 'Callback can not be empty.'
         return nil
       end
 
-
10000
      unless notification_callback.is_a? Method
+      unless notification_callback.respond_to? :call
         @logger.log Logger::ERROR, 'Invalid notification callback given.'
         return nil
       end
@@ -70,7 +78,7 @@ def add_notification_listener(notification_type, notification_callback)
     #
     # @param notification_id
     #
-    # @return [Boolean] The function returns true if found and removed, false otherwise
+    # @return [Boolean] true if found and removed, false otherwise
 
     def remove_notification_listener(notification_id)
       unless notification_id

diff --git a/spec/notification_center_spec.rb b/spec/notification_center_spec.rb
@@ -30,7 +30,11 @@
 
   before(:context) do
     class CallBack
-      def call; end
+      attr_reader :args
+
+      def call(*args)
+        @args = args
+      end
     end
 
     @callback = CallBack.new
@@ -73,6 +77,15 @@ def call; end
         expect(spy_logger).to have_received(:log).once
                                                  .with(Logger::ERROR, 'Invalid notification callback given.')
       end
+
+      it 'should log and return nil if given both callback and block' do
+        result = notification_center.add_notification_listener(Optimizely::NotificationCenter::NOTIFICATION_TYPES[:ACTIVATE],
+                                                               @callback_reference) {}
+
+        expect(result).to be_nil
+        expect(spy_logger).to have_received(:log).once
+                                                 .with(Logger::ERROR, 'Callback and block are mutually exclusive.')
+      end
     end
 
     describe 'test add notification with valid type and callback' do
@@ -484,6 +497,39 @@ def deliver_three; end
         expect(spy_logger).to_not have_received(:log)
           .with(Logger::INFO, 'delivered three.')
       end
+
+      it 'should send notifications to blocks' do
+        actual_args = []
+        notification_center.add_notification_listener(Optimizely::NotificationCenter::NOTIFICATION_TYPES[:TRACK]) do |*args|
+          actual_args = args
+        end
+
+        notification_center.send_notifications(Optimizely::NotificationCenter::NOTIFICATION_TYPES[:TRACK],
+                                               :arg1, 'arg2', arg3: 4)
+
+        expect(actual_args).to eq([:arg1, 'arg2', arg3: 4])
+      end
+
+      it 'should send notifications to lambdas' do
+        actual_args = []
+        notification_center.add_notification_listener(Optimizely::NotificationCenter::NOTIFICATION_TYPES[:TRACK],
+                                                      ->(*args) { actual_args = args })
+
+        notification_center.send_notifications(Optimizely::NotificationCenter::NOTIFICATION_TYPES[:TRACK],
+                                               :arg1, 'arg2', arg3: 4)
+
+        expect(actual_args).to eq([:arg1, 'arg2', arg3: 4])
+      end
+
+      it 'should send notifications to callables' do
+        callback = CallBack.new
+        notification_center.add_notification_listener(Optimizely::NotificationCenter::NOTIFICATION_TYPES[:TRACK], callback)
+
+        notification_center.send_notifications(Optimizely::NotificationCenter::NOTIFICATION_TYPES[:TRACK],
+                                               :arg1, 'arg2', arg3: 4)
+
+        expect(callback.args).to eq([:arg1, 'arg2', arg3: 4])
+      end
     end
 
     describe '.notification_count' do