add callback to Msgthr#add

The motivation is to allow the client to have a custom code executed,
    whenever a child is added.

    Signed-off-by: Dimid Duchovny <dimidd@gmail.com>

2 files changed, 13 insertions, 2 deletions

diff --git a/lib/msgthr.rb b/lib/msgthr.rb
index 1517f28..d88b32b 100644
--- a/lib/msgthr.rb
+++ b/lib/msgthr.rb
@@ -166,12 +166,16 @@ class Msgthr
       # but do not change existing links or loop
       if prev && !cont.parent && !cont.has_descendent(prev)
         prev.add_child(cont)
+        yield(prev, cont) if block_given?
       end
       prev = cont
     end
 
     # set parent of this message to be the last element in refs
-    prev.add_child(cur) if prev
+    if prev
+      prev.add_child(cur)
+      yield(prev, cur) if block_given?
+    end
   end
 end
 
diff --git a/test/test_msgthr.rb b/test/test_msgthr.rb
index b14e135..0f31762 100644
--- a/test/test_msgthr.rb
+++ b/test/test_msgthr.rb
@@ -6,8 +6,15 @@ require 'msgthr'
 class TestMsgthr < Test::Unit::TestCase
   def test_msgthr
     thr = Msgthr.new
+    parent_child = ''
+    # Note that C is added after B,
+    # hence it's message will be empty after adding B
+    expected_parent_child = '->B'
     thr.add('a', %w(c b), 'abc')
-    thr.add('b', %w(c), 'B')
+    thr.add('b', %w(c), 'B') do |parent, child|
+      parent_child = "#{parent.msg}->#{child.msg}"
+    end
+    assert_equal parent_child, expected_parent_child
     thr.add('c', nil, 'c')
     thr.add('D', nil, 'D')
     thr.add('d', %w(missing), 'd')