Merge pull request #157 from ar3/fix-string-vs-symbol-issue

The only and exclude options weren't working for me...
voormedia · Aug 6, 2015 · 3ed9d79 · 3ed9d79
2 parents 484df63 + f59ead8
commit 3ed9d79
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/lib/rails_erd/diagram.rb b/lib/rails_erd/diagram.rb
@@ -150,8 +150,8 @@ def callbacks
 
     def filtered_entities
       @domain.entities.reject { |entity|
-        options.exclude && entity.model && [options.exclude].flatten.include?(entity.name.to_sym) or
-        options.only && entity.model && ![options.only].flatten.include?(entity.name.to_sym) or
+        options.exclude && entity.model && [options.exclude].flatten.map(&:to_sym).include?(entity.name.to_sym) or
+        options.only && entity.model && ![options.only].flatten.map(&:to_sym).include?(entity.name.to_sym) or
         !options.inheritance && entity.specialized? or
         !options.polymorphism && entity.generalized? or
         !options.disconnected && entity.disconnected?

diff --git a/test/unit/diagram_test.rb b/test/unit/diagram_test.rb
@@ -149,6 +149,13 @@ def calls
     assert_equal [Author, Editor], retrieve_entities(:only => [:Author, :Editor]).map(&:model)
   end
 
+  test "generate should include only specified entities (With the class names as strings)" do
+    create_model "Book"
+    create_model "Author"
+    create_model "Editor"
+    assert_equal [Author, Editor], retrieve_entities(:only => ['Author', 'Editor']).map(&:model)
+  end
+
   test "generate should filter disconnected entities if disconnected is false" do
     create_model "Book", :author => :references do
       belongs_to :author