Merge pull request #145 from matugm/patch-1

Improve Attribute#from_model

Author: kerrizor

lib/rails_erd/domain/attribute.rb

@@ -14,13 +14,18 @@ def from_model(domain, model) # @private :nodoc:
           attributes.sort! if RailsERD.options[:sort]
 
           if RailsERD.options[:prepend_primary]
-            primary_key = ActiveRecord::Base.get_primary_key(model)
-            primary = attributes.detect{ |column| column.name == primary_key }
+            attributes = prepend_primary(model, attributes)
+          end
+
+          attributes
+        end
+
+        def prepend_primary(model, attributes)
+          primary_key = ActiveRecord::Base.get_primary_key(model)
+          primary = attributes.index { |column| column.name == primary_key }
 
-            if primary
-              attributes.delete(primary)
-              attributes.unshift(primary)
-            end
+          if primary
+            attributes[primary], attributes[0] = attributes[0], attributes[primary]
           end
 
           attributes

test/unit/attribute_test.rb

@@ -44,8 +44,8 @@ def create_attribute(model, name)
     assert_equal %w{id a}, Domain::Attribute.from_model(Domain.new, Foo).map(&:name)
   end
 
-  test "from_model should return attributes with PK first if sort is false and prepend_primary is true" do
-    RailsERD.options[:sort]            = false
+  test "from_model should return attributes with PK first if prepend_primary is true" do
+    RailsERD.options[:sort]            = true
     RailsERD.options[:prepend_primary] = true
 
     create_model "Foo"