Generate belongs_to association when generating a model

3 files changed, 13 insertions, 1 deletions

diff --git a/railties/lib/rails_generator/generated_attribute.rb b/railties/lib/rails_generator/generated_attribute.rb
index 25af3931de..a3d4a01142 100644
--- a/railties/lib/rails_generator/generated_attribute.rb
+++ b/railties/lib/rails_generator/generated_attribute.rb
@@ -37,6 +37,10 @@ module Rails
             ""
         end      
       end
+
+      def reference?
+        [ :references, :belongs_to ].include?(self.type)
+      end
     end
   end
 end
diff --git a/railties/lib/rails_generator/generators/components/model/templates/model.rb b/railties/lib/rails_generator/generators/components/model/templates/model.rb
index 8bf93b9e8b..6fcf393bdf 100644
--- a/railties/lib/rails_generator/generators/components/model/templates/model.rb
+++ b/railties/lib/rails_generator/generators/components/model/templates/model.rb
@@ -1,5 +1,5 @@
 class <%= class_name %> < ActiveRecord::Base
-<% attributes.select { |a| a.type.to_s == 'references' }.each do |attribute| -%>
+<% attributes.select(&:reference?).each do |attribute| -%>
   belongs_to :<%= attribute.name %>
 <% end -%>
 end
diff --git a/railties/test/generators/rails_model_generator_test.rb b/railties/test/generators/rails_model_generator_test.rb
index 67cd3aa26f..aea2abafba 100644
--- a/railties/test/generators/rails_model_generator_test.rb
+++ b/railties/test/generators/rails_model_generator_test.rb
@@ -37,4 +37,12 @@ class RailsModelGeneratorTest < GeneratorTestCase
       assert body =~ /^\s+belongs_to :supplier/, "#{body.inspect} should contain 'belongs_to :supplier'"
     end
   end
+
+  def test_model_with_belongs_to_attributes_generates_belongs_to_associations
+    run_generator('model', %w(Product name:string supplier:belongs_to))
+
+    assert_generated_model_for :product do |body|
+      assert body =~ /^\s+belongs_to :supplier/, "#{body.inspect} should contain 'belongs_to :supplier'"
+    end
+  end
 end