lib/rails_erd/domain/relationship.rb in rails-erd-0.4.0 vs lib/rails_erd/domain/relationship.rb in rails-erd-0.4.1

@@ -28,11 +28,11 @@
         def association_owner(association)
           association.options[:as] ? association.options[:as].to_s.classify : association.active_record.name
         end
         
         def association_target(association)
-          association.class_name
+          association.options[:polymorphic] ? association.class_name : association.klass.name
         end
       end
       
       extend Inspectable
       inspection_attributes :source, :destination
@@ -56,11 +56,11 @@
         @reverse_associations, @forward_associations = *unless any_habtm?(associations)
           associations.partition(&:belongs_to?)
         else
           # Many-to-many associations don't have a clearly defined direction.
           # We sort by name and use the first model as the source.
-          source = associations.first.active_record
-          associations.partition { |association| association.active_record == source }
+          source = associations.map(&:active_record).sort_by(&:name).first
+          associations.partition { |association| association.active_record != source }
         end
       
         assoc = @forward_associations.first || @reverse_associations.first
         @source = @domain.entity_by_name(self.class.send(:association_owner, assoc))
         @destination = @domain.entity_by_name(self.class.send(:association_target, assoc))