From f59ead8bdecd282f1429b5a38b06c753b6be4e37 Mon Sep 17 00:00:00 2001 From: AR3 Date: Fri, 17 Jul 2015 09:57:19 -0400 Subject: [PATCH] The only and exclude options weren't working for me, and it was because the filtering was comparing strings to symbols. --- lib/rails_erd/diagram.rb | 4 ++-- test/unit/diagram_test.rb | 7 +++++++ 2 files changed, 9 insertions(+), 2 deletions(-) diff --git a/lib/rails_erd/diagram.rb b/lib/rails_erd/diagram.rb index 6004ff35..1b4fe64b 100644 --- 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 index dd39ff13..48f0496d 100644 --- 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