diff --git a/lib/draper/base.rb b/lib/draper/base.rb index d08bb7d..2581cbb 100644 --- a/lib/draper/base.rb +++ b/lib/draper/base.rb @@ -64,15 +64,21 @@ module Draper def self.decorates_association(association_symbol, options = {}) define_method(association_symbol) do orig_association = model.send(association_symbol) - return orig_association if orig_association.nil? - if options[:with] - options[:with].decorate(orig_association) - elsif options[:polymorphic] - "#{orig_association.class}Decorator".constantize.decorate(orig_association) + + return orig_association if orig_association.nil? + + return options[:with].decorate(orig_association) if options[:with] + + if options[:polymorphic] + klass = orig_association.class + elsif model.class.respond_to?(:reflect_on_association) && model.class.reflect_on_association(association_symbol) + klass = model.class.reflect_on_association(association_symbol).klass + elsif orig_association.respond_to?(:first) + klass = orig_association.first.class else - reflection = model.class.reflect_on_association(association_symbol) - "#{reflection.klass}Decorator".constantize.decorate(orig_association, options) + klass = orig_association.class end + "#{klass}Decorator".constantize.decorate(orig_association, options) end end diff --git a/spec/draper/base_spec.rb b/spec/draper/base_spec.rb index aba9e4c..66fd6d0 100644 --- a/spec/draper/base_spec.rb +++ b/spec/draper/base_spec.rb @@ -113,22 +113,31 @@ describe Draper::Base do end context(".decorates_association") do - context "for collection associations" do + context "for ActiveModel collection associations" do before(:each){ subject.class_eval{ decorates_association :similar_products } } it "causes the association's method to return a collection of wrapped objects" do subject.similar_products.each{ |decorated| decorated.should be_instance_of(ProductDecorator) } end end - context "for a singular association" do + context "for Plain Old Ruby Object collection associations" do + before(:each){ subject.class_eval{ decorates_association :poro_similar_products } } + it "causes the association's method to return a collection of wrapped objects" do + subject.poro_similar_products.each{ |decorated| decorated.should be_instance_of(ProductDecorator) } + end + end + + context "for an ActiveModel singular association" do before(:each){ subject.class_eval{ decorates_association :previous_version } } it "causes the association's method to return a single wrapped object if the association is singular" do subject.previous_version.should be_instance_of(ProductDecorator) end + end - it "causes the association's method to return nil if the association is nil" do - source.stub(:previous_version){ nil } - subject.previous_version.should be_nil + context "for a Plain Old Ruby Object singular association" do + before(:each){ subject.class_eval{ decorates_association :poro_previous_version } } + it "causes the association's method to return a single wrapped object" do + subject.poro_previous_version.should be_instance_of(ProductDecorator) end end @@ -145,6 +154,16 @@ describe Draper::Base do subject.thing.should be_instance_of(SomeThingDecorator) end end + + context "when the association is nil" do + before(:each) do + subject.class_eval{ decorates_association :previous_version } + source.stub(:previous_version){ nil } + end + it "causes the association's method to return nil" do + subject.previous_version.should be_nil + end + end end context('.decorates_associations') do diff --git a/spec/support/samples/product.rb b/spec/support/samples/product.rb index 9a17878..a2f73ec 100644 --- a/spec/support/samples/product.rb +++ b/spec/support/samples/product.rb @@ -54,7 +54,7 @@ class Product < ActiveRecord::Base end def self.reflect_on_association(association_symbol) - OpenStruct.new(:klass => self) + association_symbol.to_s.starts_with?("poro") ? nil : OpenStruct.new(:klass => self) end def similar_products @@ -68,4 +68,13 @@ class Product < ActiveRecord::Base def thing SomeThing.new end + + def poro_similar_products + [Product.new, Product.new] + end + + def poro_previous_version + Product.new + end + end