test/unit/diagram_test.rb in rails-erd-0.4.5 vs test/unit/diagram_test.rb in rails-erd-1.0.0

@@ -2,15 +2,15 @@
 
 class DiagramTest < ActiveSupport::TestCase
   def setup
     load "rails_erd/diagram.rb"
   end
-  
+
   def teardown
     RailsERD.send :remove_const, :Diagram
   end
-  
+
   def retrieve_entities(options = {})
     klass = Class.new(Diagram)
     [].tap do |entities|
       klass.class_eval do
         each_entity do |entity, attributes|
@@ -18,11 +18,11 @@
         end
       end
       klass.create(options)
     end
   end
-  
+
   def retrieve_relationships(options = {})
     klass = Class.new(Diagram)
     [].tap do |relationships|
       klass.class_eval do
         each_relationship do |relationship|
@@ -30,11 +30,11 @@
         end
       end
       klass.create(options)
     end
   end
-  
+
   def retrieve_specializations(options = {})
     klass = Class.new(Diagram)
     [].tap do |specializations|
       klass.class_eval do
         each_specialization do |specialization|
@@ -54,11 +54,11 @@
         end
       end
       klass.create(options)
     end
   end
-  
+
   # Diagram ==================================================================
   test "domain sould return given domain" do
     domain = Object.new
     assert_same domain, Class.new(Diagram).new(domain).domain
   end
@@ -82,23 +82,23 @@
     create_simple_domain
     executed_calls = Class.new(Diagram) do
       setup do
         calls << :setup
       end
-      
+
       each_entity do
         calls << :entity
       end
-      
+
       each_relationship do
         calls << :relationship
       end
-      
+
       save do
         calls << :save
       end
-      
+
       def calls
         @calls ||= []
       end
     end.create
     assert_equal [:setup, :entity, :entity, :relationship, :save], executed_calls
@@ -121,17 +121,43 @@
         "foobar"
       end
     end.new(Domain.generate)
     assert_equal "foobar", diagram.create
   end
-  
+
   # Entity filtering =========================================================
   test "generate should yield entities" do
     create_model "Foo"
     assert_equal [Foo], retrieve_entities.map(&:model)
   end
 
+  test "generate should filter excluded entity" do
+    create_model "Book"
+    create_model "Author"
+    assert_equal [Book], retrieve_entities(:exclude => :Author).map(&:model)
+  end
+
+  test "generate should filter excluded entities" do
+    create_model "Book"
+    create_model "Author"
+    create_model "Editor"
+    assert_equal [Book], retrieve_entities(:exclude => [:Author, :Editor]).map(&:model)
+  end
+
+  test "generate should include only specified entity" do
+    create_model "Book"
+    create_model "Author"
+    assert_equal [Book], retrieve_entities(:only => [:Book]).map(&:model)
+  end
+
+  test "generate should include only specified entities" 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
     end
     create_model "Author"
@@ -147,22 +173,22 @@
   test "generate should filter specialized entities" do
     create_model "Foo", :type => :string
     Object.const_set :SpecialFoo, Class.new(Foo)
     assert_equal [Foo], retrieve_entities.map(&:model)
   end
-  
+
   test "generate should yield specialized entities if inheritance is true" do
     create_model "Foo", :type => :string
     Object.const_set :SpecialFoo, Class.new(Foo)
     assert_equal [Foo, SpecialFoo], retrieve_entities(:inheritance => true).map(&:model)
   end
 
   test "generate should yield specialized entities with distinct tables" do
     create_model "Foo"
     Object.const_set :SpecialFoo, Class.new(Foo)
     SpecialFoo.class_eval do
-      set_table_name "special_foo"
+      self.table_name = "special_foo"
     end
     create_table "special_foo", {}, true
     assert_equal [Foo, SpecialFoo], retrieve_entities.map(&:model)
   end
 
@@ -171,19 +197,19 @@
     create_model "Galleon" do
       has_many :cannons, :as => :defensible
     end
     assert_equal ["Cannon", "Galleon"], retrieve_entities.map(&:name)
   end
-  
+
   test "generate should yield generalized entities if polymorphism is true" do
     create_model "Cannon"
     create_model "Galleon" do
       has_many :cannons, :as => :defensible
     end
     assert_equal ["Cannon", "Defensible", "Galleon"], retrieve_entities(:polymorphism => true).map(&:name)
   end
-  
+
   # Relationship filtering ===================================================
   test "generate should yield relationships" do
     create_simple_domain
     assert_equal 1, retrieve_relationships.length
   end
@@ -200,11 +226,11 @@
     create_model "Baz", :foo => :references do
       belongs_to :foo
     end
     assert_equal [false, false, true], retrieve_relationships(:indirect => true).map(&:indirect?)
   end
-  
+
   test "generate should filter indirect relationships if indirect is false" do
     create_model "Foo" do
       has_many :bazs
       has_many :bars
     end
@@ -225,11 +251,11 @@
     SpecialBar.class_eval do
       has_many :foos
     end
     assert_equal 1, retrieve_relationships.length
   end
-  
+
   test "generate should yield relationships to specialized entities" do
     create_model "Foo", :type => :string, :bar => :references
     Object.const_set :SpecialFoo, Class.new(Foo)
     create_model "Bar" do
       has_many :special_foos
@@ -241,11 +267,11 @@
   test "generate should not yield specializations" do
     create_specialization
     create_generalization
     assert_equal [], retrieve_specializations
   end
-  
+
   test "generate should yield specializations but not generalizations if inheritance is true" do
     create_specialization
     create_generalization
     assert_equal ["Beer"], retrieve_specializations(:inheritance => true).map { |s| s.specialized.name }
   end
@@ -303,10 +329,10 @@
     end
     create_model "Author"
     assert_equal %w{created_at title},
       retrieve_attribute_lists(:attributes => [:content, :timestamps])[Book].map(&:name)
   end
-  
+
   test "generate should yield no attributes for specialized entities" do
     create_model "Beverage", :type => :string, :name => :string, :distillery => :string, :age => :integer
     Object.const_set :Whisky, Class.new(Beverage)
     assert_equal [], retrieve_attribute_lists(:inheritance => true)[Whisky].map(&:name)
   end