Support decorating Mongoid::Criteria

Gemfile

@@ -6,3 +6,4 @@ gem 'rails', '~> 3.2.0'
 gem 'rspec-rails'
 gem 'capybara', '>= 2'
 gem 'sqlite3'
+gem 'mongoid'

lib/active_decorator.rb

@@ -1,3 +1,4 @@
 require 'active_decorator/version'
 require 'active_decorator/decorator'
 require 'active_decorator/railtie'
+require 'active_decorator/orm/mongoid'

lib/active_decorator/decorator.rb

@@ -25,6 +25,8 @@ def to_a_with_decorator
           end
           alias_method_chain :to_a, :decorator
         end
+      elsif defined?(Mongoid) && obj.is_a?(Mongoid::Criteria)
+        obj.extend ORM::Mongoid
       else
         d = decorator_for obj.class
         return obj unless d

lib/active_decorator/orm/mongoid.rb

@@ -0,0 +1,31 @@
+module ActiveDecorator
+  module ORM
+    module Mongoid
+      def first
+        decorate super
+      end
+
+      def last
+        decorate super
+      end
+
+      def each(&block)
+        super do |document|
+          yield decorate(document)
+        end
+      end
+
+      def map(&block)
+        super do |document|
+          yield decorate(document)
+        end
+      end
+
+      private
+
+      def decorate(object)
+        ActiveDecorator::Decorator.instance.decorate(object)
+      end
+    end
+  end
+end

spec/decorator_spec.rb

@@ -0,0 +1,25 @@
+require 'spec_helper'
+
+describe ActiveDecorator::Decorator do
+  subject do
+    ActiveDecorator::Decorator.instance
+  end
+
+  describe '#decorate' do
+    context 'Mongoid::Criteria' do
+      let(:criteria) {Mongoid::Criteria.new(MongoidDummy)}
+
+      it 'is extended with ORM::Mongoid' do
+        criteria.should_receive(:extend).with(ActiveDecorator::ORM::Mongoid)
+        subject.decorate(criteria)
+      end
+    end
+  end
+end
+
+class MongoidDummy
+  include Mongoid::Document
+end
+
+module MongoidDummyDecorator
+end

spec/orm/mongoid_spec.rb

@@ -0,0 +1,40 @@
+require 'spec_helper'
+
+describe ActiveDecorator::ORM::Mongoid do
+  subject do
+    MongoidDummy.create
+    MongoidDummy.all.extend(ActiveDecorator::ORM::Mongoid)
+  end
+
+  it 'decorates #first' do
+    subject.first.thing.should be_true
+  end
+
+  it 'decorates #last' do
+    subject.last.thing.should be_true
+  end
+
+  it 'decorates #each' do
+    subject.each do |document|
+      document.thing.should be_true
+    end
+  end
+
+  it 'decorates #map' do
+    subject.map(&:thing).should == [true]
+  end
+end
+
+class MongoidDummy
+  include Mongoid::Document
+
+  def thing
+    false
+  end
+end
+
+module MongoidDummyDecorator
+  def thing
+    true
+  end
+end

spec/spec_helper.rb

@@ -6,6 +6,7 @@
 # needs to load the app before loading rspec/rails => capybara
 require 'fake_app/fake_app'
 require 'rspec/rails'
+require 'mongoid'
 # Requires supporting files with custom matchers and macros, etc,
 # in ./support/ and its subdirectories.
 Dir["#{File.dirname(__FILE__)}/support/**/*.rb"].each {|f| require f}
@@ -18,5 +19,10 @@
     Book.delete_all
     Author.delete_all
     Movie.delete_all
+    Mongoid::Config.purge!
   end
 end
+
+Mongoid.configure do |config|
+  config.connect_to('active_decorator_test')
+end