adding save and initialize callbacks

lib/mini_mongo/concerns/callbacks.rb

@@ -8,7 +8,8 @@ module Callbacks
       included do
         extend ActiveModel::Callbacks  
 
-        define_model_callbacks :insert, :update, :remove
+        define_model_callbacks :insert, :update, :save, :remove, :only => [:after, :before]
+        define_model_callbacks :initialize, :only => [:after]
       end
     end
   end

lib/mini_mongo/concerns/persistance.rb

@@ -34,16 +34,17 @@ def save!(*args)
       def insert(options = {})
         raise AlreadyInsertedError, "document has already been inserted" if persisted?
 
-        response = run_callbacks :insert do
-          validate!
-          response = collection.insert(document.to_hash, options)
-          raise MiniMongo::InsertError, "blank _id: #{response.inspect}" if response.blank?
-          document.dot_set("_id", response) if document["_id"].blank?
-          @persisted = true
-          clear_snapshot
-          response
+        run_callbacks :save do
+          run_callbacks :insert do
+            validate!
+            response = collection.insert(document.to_hash, options)
+            raise MiniMongo::InsertError, "blank _id: #{response.inspect}" if response.blank?
+            document.dot_set("_id", response) if document["_id"].blank?
+            @persisted = true
+            clear_snapshot
+            response
+          end
         end
-        response
       rescue Mongo::OperationFailure => e
         if e.message.to_s =~ /^11000\:/
           raise MiniMongo::DuplicateKeyError, e.message
@@ -61,28 +62,30 @@ def update(options = {})
 
         raise NotInsertedError, "document must be inserted before being updated" unless persisted?
 
-        run_callbacks :update do
-          validate!
-          only_if_current = options.delete(:only_if_current)
-          options[:safe] = true if !options[:safe] && only_if_current
-          selector = self.class.build_update_selector(self.to_oid, changes, only_if_current)
-          updates = self.class.build_update_hash(changes)
-
-          if options.delete(:find_and_modify) == true
-            response = self.collection.find_and_modify(query: selector, update: updates, new: true)
-            reload(response)
-          else
-            response = collection.update(selector, updates, options)
-            if !response.is_a?(Hash) || (response["updatedExisting"] && response["n"] == 1)
-              clear_snapshot
-              true
+        run_callbacks :save do
+          run_callbacks :update do
+            validate!
+            only_if_current = options.delete(:only_if_current)
+            options[:safe] = true if !options[:safe] && only_if_current
+            selector = self.class.build_update_selector(self.to_oid, changes, only_if_current)
+            updates = self.class.build_update_hash(changes)
+
+            if options.delete(:find_and_modify) == true
+              response = self.collection.find_and_modify(query: selector, update: updates, new: true)
+              reload(response)
             else
-              if only_if_current
-                raise MiniMongo::StaleUpdateError, response.inspect
+              response = collection.update(selector, updates, options)
+              if !response.is_a?(Hash) || (response["updatedExisting"] && response["n"] == 1)
+                clear_snapshot
+                true
               else
-                raise MiniMongo::UpdateError, response.inspect
+                if only_if_current
+                  raise MiniMongo::StaleUpdateError, response.inspect
+                else
+                  raise MiniMongo::UpdateError, response.inspect
+                end
               end
-            end
+            end            
           end
         end
       rescue Mongo::OperationFailure => e

lib/mini_mongo/document.rb

@@ -10,7 +10,7 @@ module MiniMongo::Document
 
   def initialize(document = {}, persisted = false)
     set_document(prepare_document(document, persisted))
-    @persisted = persisted
+    run_callbacks :initialize unless @persisted = persisted
   end
 
   def document

test/test_callbacks.rb

@@ -0,0 +1,129 @@
+require_relative 'helper'
+
+class TestCallbacks < MiniTest::Unit::TestCase
+
+  class Person
+    include MiniMongo::Document
+
+    attr_accessor :initialize_callback
+    attr_accessor :insert_callback
+    attr_accessor :update_callback
+    attr_accessor :save_callback
+    attr_accessor :remove_callback
+
+    set_callback :initialize do |object|
+      object.initialize_callback = true
+    end
+
+    set_callback :insert do |object|
+      object.insert_callback = true
+    end
+
+    set_callback :save do |object|
+      object.save_callback = true
+    end
+
+    set_callback :update do |object|
+      object.update_callback = true
+    end
+
+    set_callback :remove do |object|
+      object.remove_callback = true
+    end
+  end
+
+  def test_callbacks_on_initialize
+    person = Person.new
+    assert person.initialize_callback
+    person.insert!
+    person = Person.first
+    refute person.initialize_callback
+  end
+
+  def test_callbacks_on_insert
+    person = Person.new
+    person.stubs(dirty?: true)
+
+    refute person.insert_callback
+    refute person.save_callback
+    refute person.update_callback
+    refute person.remove_callback
+
+    person.insert!
+
+    assert person.insert_callback
+    refute person.update_callback
+    assert person.save_callback
+    refute person.remove_callback
+  end
+
+  def test_callbacks_on_update
+    Person.new.insert!
+    person = Person.first
+    person.stubs(dirty?: true)
+
+    refute person.insert_callback
+    refute person.save_callback
+    refute person.update_callback
+    refute person.remove_callback
+
+    person.update!
+
+    refute person.insert_callback
+    assert person.update_callback
+    assert person.save_callback
+    refute person.remove_callback
+  end
+
+  def test_callbacks_on_save_when_new
+    person = Person.new
+    person.stubs(dirty?: true)
+
+    refute person.insert_callback
+    refute person.save_callback
+    refute person.update_callback
+    refute person.remove_callback
+
+    person.save!
+
+    assert person.insert_callback
+    refute person.update_callback
+    assert person.save_callback
+    refute person.remove_callback
+  end
+
+  def test_callbacks_on_save_when_update
+    Person.new.insert!
+    person = Person.first
+    person.stubs(dirty?: true)
+
+    refute person.insert_callback
+    refute person.save_callback
+    refute person.update_callback
+    refute person.remove_callback
+
+    person.save!
+
+    refute person.insert_callback
+    assert person.update_callback
+    assert person.save_callback
+    refute person.remove_callback
+  end
+
+  def test_callbacks_on_remove
+    Person.new.insert!
+    person = Person.first
+
+    refute person.insert_callback
+    refute person.save_callback
+    refute person.update_callback
+    refute person.remove_callback
+
+    person.remove!
+
+    refute person.insert_callback
+    refute person.update_callback
+    refute person.save_callback
+    assert person.remove_callback
+  end
+end