From 108d4baaafb5990872c9eea979a2ce1f40bb0f31 Mon Sep 17 00:00:00 2001
From: Simon COURTOIS <scourtois@cubyx.fr>
Date: Thu, 8 Mar 2012 14:21:26 +0100
Subject: [PATCH] Adding a :polymorphic option to decorates_association

---
 lib/draper/base.rb                           | 2 ++
 spec/draper/base_spec.rb                     | 7 +++++++
 spec/spec_helper.rb                          | 2 ++
 spec/support/samples/product.rb              | 4 ++++
 spec/support/samples/some_thing.rb           | 2 ++
 spec/support/samples/some_thing_decorator.rb | 3 +++
 6 files changed, 20 insertions(+)
 create mode 100644 spec/support/samples/some_thing.rb
 create mode 100644 spec/support/samples/some_thing_decorator.rb

diff --git a/lib/draper/base.rb b/lib/draper/base.rb
index 9ae4d4a4..48a3316a 100644
--- a/lib/draper/base.rb
+++ b/lib/draper/base.rb
@@ -67,6 +67,8 @@ def self.decorates_association(association_symbol, options = {})
         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)
         else
           reflection = model.class.reflect_on_association(association_symbol)
           "#{reflection.klass}Decorator".constantize.decorate(orig_association)
diff --git a/spec/draper/base_spec.rb b/spec/draper/base_spec.rb
index 6a0fe4b9..9bc492d6 100644
--- a/spec/draper/base_spec.rb
+++ b/spec/draper/base_spec.rb
@@ -138,6 +138,13 @@ class CustomDecorator < Draper::Base
         subject.previous_version.should be_instance_of(SpecificProductDecorator)
       end
     end
+
+    context "for a polymorphic association" do
+      before(:each){ subject.class_eval{ decorates_association :thing, :polymorphic => true } }
+      it "causes the association to be decorated with the right decorator" do
+        subject.thing.should be_instance_of(SomeThingDecorator)
+      end
+    end
   end
 
   context('.decorates_associations') do
diff --git a/spec/spec_helper.rb b/spec/spec_helper.rb
index 8ea84da2..6b4e82a5 100644
--- a/spec/spec_helper.rb
+++ b/spec/spec_helper.rb
@@ -15,5 +15,7 @@
 require './spec/support/samples/product.rb'
 require './spec/support/samples/product_decorator.rb'
 require './spec/support/samples/specific_product_decorator.rb'
+require './spec/support/samples/some_thing.rb'
+require './spec/support/samples/some_thing_decorator.rb'
 require './spec/support/samples/widget.rb'
 require './spec/support/samples/widget_decorator.rb'
diff --git a/spec/support/samples/product.rb b/spec/support/samples/product.rb
index 92f04252..9a17878e 100644
--- a/spec/support/samples/product.rb
+++ b/spec/support/samples/product.rb
@@ -64,4 +64,8 @@ def similar_products
   def previous_version
     Product.new
   end
+
+  def thing
+    SomeThing.new
+  end
 end
diff --git a/spec/support/samples/some_thing.rb b/spec/support/samples/some_thing.rb
new file mode 100644
index 00000000..57309d30
--- /dev/null
+++ b/spec/support/samples/some_thing.rb
@@ -0,0 +1,2 @@
+class SomeThing < Product
+end
diff --git a/spec/support/samples/some_thing_decorator.rb b/spec/support/samples/some_thing_decorator.rb
new file mode 100644
index 00000000..07844ada
--- /dev/null
+++ b/spec/support/samples/some_thing_decorator.rb
@@ -0,0 +1,3 @@
+class SomeThingDecorator < Draper::Base
+  decorates :some_thing
+end