rails-erd 0.3.0 → 0.4.0

data/lib/rails_erd/railtie.rb

@@ -1,7 +1,7 @@
 module RailsERD
   # Rails ERD integrates with Rails 3. If you add it to your +Gemfile+, you
   # will gain a Rake task called +erd+, which you can use to generate diagrams
-  # of your domain model. See the README.rdoc file for more information.
+  # of your domain model.
   class Railtie < Rails::Railtie
     rake_tasks do
       load "rails_erd/tasks.rake"

data/rails-erd.gemspec

@@ -5,16 +5,16 @@
 
 Gem::Specification.new do |s|
   s.name = %q{rails-erd}
-  s.version = "0.3.0"
+  s.version = "0.4.0"
 
   s.required_rubygems_version = Gem::Requirement.new(">= 0") if s.respond_to? :required_rubygems_version=
   s.authors = ["Rolf Timmermans"]
-  s.date = %q{2010-10-03}
+  s.date = %q{2010-10-12}
   s.description = %q{Automatically generate an entity-relationship diagram (ERD) for your Rails models.}
   s.email = %q{r.timmermans@voormedia.com}
   s.extra_rdoc_files = [
     "LICENSE",
-     "README.rdoc"
+     "README.md"
   ]
   s.files = [
     ".gitignore",
@@ -22,20 +22,21 @@ Gem::Specification.new do |s|
      "Gemfile",
      "Gemfile.lock",
      "LICENSE",
-     "README.rdoc",
+     "README.md",
      "Rakefile",
      "VERSION",
      "lib/rails-erd.rb",
      "lib/rails_erd.rb",
-     "lib/rails_erd/attribute.rb",
      "lib/rails_erd/diagram.rb",
      "lib/rails_erd/diagram/graphviz.rb",
      "lib/rails_erd/diagram/templates/node.erb",
      "lib/rails_erd/domain.rb",
-     "lib/rails_erd/entity.rb",
+     "lib/rails_erd/domain/attribute.rb",
+     "lib/rails_erd/domain/entity.rb",
+     "lib/rails_erd/domain/relationship.rb",
+     "lib/rails_erd/domain/relationship/cardinality.rb",
+     "lib/rails_erd/domain/specialization.rb",
      "lib/rails_erd/railtie.rb",
-     "lib/rails_erd/relationship.rb",
-     "lib/rails_erd/relationship/cardinality.rb",
      "lib/rails_erd/tasks.rake",
      "rails-erd.gemspec",
      "test/test_helper.rb",
@@ -46,7 +47,8 @@ Gem::Specification.new do |s|
      "test/unit/entity_test.rb",
      "test/unit/graphviz_test.rb",
      "test/unit/rake_task_test.rb",
-     "test/unit/relationship_test.rb"
+     "test/unit/relationship_test.rb",
+     "test/unit/specialization_test.rb"
   ]
   s.homepage = %q{http://rails-erd.rubyforge.org/}
   s.rdoc_options = ["--charset=UTF-8"]
@@ -63,7 +65,8 @@ Gem::Specification.new do |s|
      "test/unit/entity_test.rb",
      "test/unit/graphviz_test.rb",
      "test/unit/rake_task_test.rb",
-     "test/unit/relationship_test.rb"
+     "test/unit/relationship_test.rb",
+     "test/unit/specialization_test.rb"
   ]
 
   if s.respond_to? :specification_version then
@@ -71,20 +74,20 @@ Gem::Specification.new do |s|
     s.specification_version = 3
 
     if Gem::Version.new(Gem::VERSION) >= Gem::Version.new('1.2.0') then
-      s.add_runtime_dependency(%q<activerecord>, ["~> 3.0.0"])
+      s.add_runtime_dependency(%q<activerecord>, ["~> 3.0"])
       s.add_runtime_dependency(%q<activesupport>, ["~> 3.0"])
-      s.add_runtime_dependency(%q<ruby-graphviz>, ["~> 0.9.17"])
+      s.add_runtime_dependency(%q<ruby-graphviz>, ["~> 0.9.18"])
       s.add_development_dependency(%q<sqlite3-ruby>, [">= 0"])
     else
-      s.add_dependency(%q<activerecord>, ["~> 3.0.0"])
+      s.add_dependency(%q<activerecord>, ["~> 3.0"])
       s.add_dependency(%q<activesupport>, ["~> 3.0"])
-      s.add_dependency(%q<ruby-graphviz>, ["~> 0.9.17"])
+      s.add_dependency(%q<ruby-graphviz>, ["~> 0.9.18"])
       s.add_dependency(%q<sqlite3-ruby>, [">= 0"])
     end
   else
-    s.add_dependency(%q<activerecord>, ["~> 3.0.0"])
+    s.add_dependency(%q<activerecord>, ["~> 3.0"])
     s.add_dependency(%q<activesupport>, ["~> 3.0"])
-    s.add_dependency(%q<ruby-graphviz>, ["~> 0.9.17"])
+    s.add_dependency(%q<ruby-graphviz>, ["~> 0.9.18"])
     s.add_dependency(%q<sqlite3-ruby>, [">= 0"])
   end
 end

data/test/test_helper.rb

@@ -1,7 +1,7 @@
 require "rubygems"
-require "bundler"
-Bundler.require
+require "bundler/setup"
 
+require "active_record"
 require "test/unit"
 require "rails_erd/domain"
 
@@ -33,10 +33,14 @@ class ActiveSupport::TestCase
     end
   end
 
-  def create_model(name, columns = {}, &block)
-    klass = Object.const_set name.to_sym, Class.new(ActiveRecord::Base)
+  def create_model(name, *args, &block)
+    superklass = args.first.kind_of?(Class) ? args.shift : ActiveRecord::Base
+    columns = args.first || {}
+    klass = Object.const_set name.to_sym, Class.new(superklass)
     klass.class_eval(&block) if block_given?
-    create_table Object.const_get(name.to_sym).table_name, columns, Object.const_get(name.to_sym).primary_key rescue nil
+    if superklass == ActiveRecord::Base
+      create_table Object.const_get(name.to_sym).table_name, columns, Object.const_get(name.to_sym).primary_key rescue nil
+    end
   end
     
   def create_models(*names)
@@ -90,6 +94,18 @@ class ActiveSupport::TestCase
     create_table "manies_mores", :many_id => :integer, :more_id => :integer
   end
   
+  def create_specialization
+    create_model "Beverage", :type => :string
+    Object.const_set :Beer, Class.new(Beverage)
+  end
+  
+  def create_generalization
+    create_model "Cannon"
+    create_model "Galleon" do
+      has_many :cannons, :as => :defensible
+    end
+  end
+
   private
   
   def reset_domain

data/test/unit/attribute_test.rb

@@ -20,28 +20,28 @@ class AttributeTest < ActiveSupport::TestCase
   end
   
   def create_attribute(model, name)
-    Attribute.new(Domain.generate, model, model.arel_table.attributes[name].column)
+    Domain::Attribute.new(Domain.generate, model, model.columns_hash[name])
   end
   
   # Attribute ================================================================
   test "column should return database column" do
     create_model "Foo", :my_column => :string
-    assert_equal Foo.arel_table.attributes["my_column"].column,
-      Attribute.from_model(Domain.new, Foo).reject(&:primary_key?).first.column
+    assert_equal Foo.columns_hash["my_column"],
+      Domain::Attribute.from_model(Domain.new, Foo).reject(&:primary_key?).first.column
   end
   
   test "spaceship should sort attributes by name" do
     create_model "Foo", :a => :string, :b => :string, :c => :string
-    a = Attribute.new(Domain.new, Foo, Foo.arel_table.attributes["a"].column)
-    b = Attribute.new(Domain.new, Foo, Foo.arel_table.attributes["b"].column)
-    c = Attribute.new(Domain.new, Foo, Foo.arel_table.attributes["c"].column)
+    a = create_attribute(Foo, "a")
+    b = create_attribute(Foo, "b")
+    c = create_attribute(Foo, "c")
     assert_equal [a, b, c], [c, a, b].sort
   end
   
   test "inspect should show column" do
     create_model "Foo", :my_column => :string
-    assert_match %r{#<RailsERD::Attribute:.* @column="my_column" @type=:string>},
-      Attribute.new(Domain.new, Foo, Foo.arel_table.attributes["my_column"].column).inspect
+    assert_match %r{#<RailsERD::Domain::Attribute:.* @name="my_column" @type=:string>},
+      Domain::Attribute.new(Domain.new, Foo, Foo.columns_hash["my_column"]).inspect
   end
   
   test "type should return attribute type" do
@@ -112,6 +112,33 @@ class AttributeTest < ActiveSupport::TestCase
       create_attribute(Foo, "created_on"), create_attribute(Foo, "updated_on")].collect(&:timestamp?)
   end
   
+  test "inheritance should return false by default" do
+    create_model "Foo", :type => :string, :alternative => :string do
+      set_inheritance_column :alternative
+    end
+    assert_equal false, create_attribute(Foo, "type").inheritance?
+  end
+  
+  test "inheritance should return if this column is used for single table inheritance" do
+    create_model "Foo", :type => :string, :alternative => :string do
+      set_inheritance_column :alternative
+    end
+    assert_equal true, create_attribute(Foo, "alternative").inheritance?
+  end
+  
+  test "content should return true by default" do
+    create_model "Foo", :my_first_column => :string
+    assert_equal true, create_attribute(Foo, "my_first_column").content?
+  end
+
+  test "content should return false for primary keys, foreign keys, timestamps and inheritance columns" do
+    create_model "Book", :type => :string, :created_at => :datetime, :case => :references do
+      belongs_to :case
+    end
+    create_model "Case"
+    assert_equal [false] * 4, %w{id type created_at case_id}.map { |a| create_attribute(Book, a).content? }
+  end
+  
   # Type descriptions ========================================================
   test "type_description should return short type description" do
     create_model "Foo", :a => :binary

data/test/unit/cardinality_test.rb

@@ -2,115 +2,121 @@ require File.expand_path("../test_helper", File.dirname(__FILE__))
 
 class CardinalityTest < ActiveSupport::TestCase
   def setup
-    @n = Relationship::Cardinality::N
+    @n = Domain::Relationship::Cardinality::N
   end
   
+  # Cardinality ==============================================================
+  test "inspect should show source and destination ranges" do
+    assert_match %r{#<RailsERD::Domain::Relationship::Cardinality:.* @source_range=1\.\.1 @destination_range=1\.\.Infinity>},
+      Domain::Relationship::Cardinality.new(1, 1..@n).inspect
+  end
+
   # Cardinality construction =================================================
   test "new should return cardinality object" do
-    assert_kind_of Relationship::Cardinality, Relationship::Cardinality.new(1, 1..@n)
+    assert_kind_of Domain::Relationship::Cardinality, Domain::Relationship::Cardinality.new(1, 1..@n)
   end
   
   # Cardinality properties ===================================================
   test "source_optional should return true if source range starts at zero" do
-    assert_equal true, Relationship::Cardinality.new(0..1, 1).source_optional?
+    assert_equal true, Domain::Relationship::Cardinality.new(0..1, 1).source_optional?
   end
 
   test "source_optional should return false if source range starts at one or more" do
-    assert_equal false, Relationship::Cardinality.new(1..2, 0..1).source_optional?
+    assert_equal false, Domain::Relationship::Cardinality.new(1..2, 0..1).source_optional?
   end
 
   test "destination_optional should return true if destination range starts at zero" do
-    assert_equal true, Relationship::Cardinality.new(1, 0..1).destination_optional?
+    assert_equal true, Domain::Relationship::Cardinality.new(1, 0..1).destination_optional?
   end
 
   test "destination_optional should return false if destination range starts at one or more" do
-    assert_equal false, Relationship::Cardinality.new(0..1, 1..2).destination_optional?
+    assert_equal false, Domain::Relationship::Cardinality.new(0..1, 1..2).destination_optional?
   end
   
   test "inverse should return inverse cardinality" do
-    assert_equal Relationship::Cardinality.new(23..45, 0..15), Relationship::Cardinality.new(0..15, 23..45).inverse
+    assert_equal Domain::Relationship::Cardinality.new(23..45, 0..15), Domain::Relationship::Cardinality.new(0..15, 23..45).inverse
   end
   
   # Cardinality equality =====================================================
   test "cardinalities are equal if they have the same boundaries" do
-    assert_equal Relationship::Cardinality.new(1, 1..Relationship::Cardinality::N),
-      Relationship::Cardinality.new(1, 1..Relationship::Cardinality::N)
+    assert_equal Domain::Relationship::Cardinality.new(1, 1..Domain::Relationship::Cardinality::N),
+      Domain::Relationship::Cardinality.new(1, 1..Domain::Relationship::Cardinality::N)
   end
 
   test "cardinalities are not equal if they have a different source range" do
-    assert_not_equal Relationship::Cardinality.new(0..1, 1..Relationship::Cardinality::N),
-      Relationship::Cardinality.new(1..1, 1..Relationship::Cardinality::N)
+    assert_not_equal Domain::Relationship::Cardinality.new(0..1, 1..Domain::Relationship::Cardinality::N),
+      Domain::Relationship::Cardinality.new(1..1, 1..Domain::Relationship::Cardinality::N)
   end
 
   test "cardinalities are not equal if they have a different destination range" do
-    assert_not_equal Relationship::Cardinality.new(0..1, 1..Relationship::Cardinality::N),
-      Relationship::Cardinality.new(0..1, 2..Relationship::Cardinality::N)
+    assert_not_equal Domain::Relationship::Cardinality.new(0..1, 1..Domain::Relationship::Cardinality::N),
+      Domain::Relationship::Cardinality.new(0..1, 2..Domain::Relationship::Cardinality::N)
   end
   
   # Cardinal names ===========================================================
   test "one_to_one should return true if source and destination are exactly one" do
-    assert_equal true, Relationship::Cardinality.new(1, 1).one_to_one?
+    assert_equal true, Domain::Relationship::Cardinality.new(1, 1).one_to_one?
   end
 
   test "one_to_one should return true if source and destination range are less than or equal to one" do
-    assert_equal true, Relationship::Cardinality.new(0..1, 0..1).one_to_one?
+    assert_equal true, Domain::Relationship::Cardinality.new(0..1, 0..1).one_to_one?
   end
 
   test "one_to_one should return false if source range upper limit is more than one" do
-    assert_equal false, Relationship::Cardinality.new(0..15, 0..1).one_to_one?
+    assert_equal false, Domain::Relationship::Cardinality.new(0..15, 0..1).one_to_one?
   end
   
   test "one_to_one should return false if destination range upper limit is more than one" do
-    assert_equal false, Relationship::Cardinality.new(0..1, 0..15).one_to_one?
+    assert_equal false, Domain::Relationship::Cardinality.new(0..1, 0..15).one_to_one?
   end
   
   test "one_to_many should return true if source is exactly one and destination is higher than one" do
-    assert_equal true, Relationship::Cardinality.new(1, 15).one_to_many?
+    assert_equal true, Domain::Relationship::Cardinality.new(1, 15).one_to_many?
   end
 
   test "one_to_many should return true if source is less than or equal to one and destination is higher than one" do
-    assert_equal true, Relationship::Cardinality.new(0..1, 0..15).one_to_many?
+    assert_equal true, Domain::Relationship::Cardinality.new(0..1, 0..15).one_to_many?
   end
 
   test "one_to_many should return false if source range upper limit is more than one" do
-    assert_equal false, Relationship::Cardinality.new(0..15, 0..15).one_to_many?
+    assert_equal false, Domain::Relationship::Cardinality.new(0..15, 0..15).one_to_many?
   end
   
   test "one_to_many should return false if destination range upper limit is one" do
-    assert_equal false, Relationship::Cardinality.new(0..1, 1).one_to_many?
+    assert_equal false, Domain::Relationship::Cardinality.new(0..1, 1).one_to_many?
   end
 
   test "many_to_many should return true if source and destination are higher than one" do
-    assert_equal true, Relationship::Cardinality.new(15, 15).many_to_many?
+    assert_equal true, Domain::Relationship::Cardinality.new(15, 15).many_to_many?
   end
 
   test "many_to_many should return true if source and destination upper limits are higher than one" do
-    assert_equal true, Relationship::Cardinality.new(0..15, 0..15).many_to_many?
+    assert_equal true, Domain::Relationship::Cardinality.new(0..15, 0..15).many_to_many?
   end
 
   test "many_to_many should return false if source range upper limit is is one" do
-    assert_equal false, Relationship::Cardinality.new(1, 0..15).many_to_many?
+    assert_equal false, Domain::Relationship::Cardinality.new(1, 0..15).many_to_many?
   end
   
   test "many_to_many should return false if destination range upper limit is one" do
-    assert_equal false, Relationship::Cardinality.new(0..1, 1).many_to_many?
+    assert_equal false, Domain::Relationship::Cardinality.new(0..1, 1).many_to_many?
   end
   
   test "inverse of one_to_many should be many_to_one" do
-    assert_equal true, Relationship::Cardinality.new(0..1, 0..@n).inverse.many_to_one?
+    assert_equal true, Domain::Relationship::Cardinality.new(0..1, 0..@n).inverse.many_to_one?
   end
 
   # Cardinality order ========================================================
   test "cardinalities should be sorted in order of maniness" do
-    card1 = Relationship::Cardinality.new(0..1, 1)
-    card2 = Relationship::Cardinality.new(1, 1)
-    card3 = Relationship::Cardinality.new(0..1, 1..3)
-    card4 = Relationship::Cardinality.new(1, 1..2)
-    card5 = Relationship::Cardinality.new(1, 1..@n)
-    card6 = Relationship::Cardinality.new(1..5, 1..3)
-    card7 = Relationship::Cardinality.new(1..2, 1..15)
-    card8 = Relationship::Cardinality.new(1..15, 1..@n)
-    card9 = Relationship::Cardinality.new(1..@n, 1..@n)
+    card1 = Domain::Relationship::Cardinality.new(0..1, 1)
+    card2 = Domain::Relationship::Cardinality.new(1, 1)
+    card3 = Domain::Relationship::Cardinality.new(0..1, 1..3)
+    card4 = Domain::Relationship::Cardinality.new(1, 1..2)
+    card5 = Domain::Relationship::Cardinality.new(1, 1..@n)
+    card6 = Domain::Relationship::Cardinality.new(1..5, 1..3)
+    card7 = Domain::Relationship::Cardinality.new(1..2, 1..15)
+    card8 = Domain::Relationship::Cardinality.new(1..15, 1..@n)
+    card9 = Domain::Relationship::Cardinality.new(1..@n, 1..@n)
     assert_equal [card1, card2, card3, card4, card5, card6, card7, card8, card9],
       [card9, card5, card8, card2, card4, card7, card1, card6, card3].sort
   end

data/test/unit/diagram_test.rb

@@ -9,35 +9,47 @@ class DiagramTest < ActiveSupport::TestCase
     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|
+          entities << entity
+        end
+      end
+      klass.create(options)
+    end
+  end
+  
   def retrieve_relationships(options = {})
     klass = Class.new(Diagram)
     [].tap do |relationships|
       klass.class_eval do
-        define_method :process_relationship do |relationship|
+        each_relationship do |relationship|
           relationships << relationship
         end
       end
       klass.create(options)
     end
   end
-
-  def retrieve_entities(options = {})
+  
+  def retrieve_specializations(options = {})
     klass = Class.new(Diagram)
-    [].tap do |entities|
+    [].tap do |specializations|
       klass.class_eval do
-        define_method :process_entity do |entity, attributes|
-          entities << entity
+        each_specialization do |specialization|
+          specializations << specialization
         end
       end
       klass.create(options)
     end
   end
-  
+
   def retrieve_attribute_lists(options = {})
     klass = Class.new(Diagram)
     {}.tap do |attribute_lists|
       klass.class_eval do
-        define_method :process_entity do |entity, attributes|
+        each_entity do |entity, attributes|
           attribute_lists[entity.model] = attributes
         end
       end
@@ -46,10 +58,56 @@ class DiagramTest < ActiveSupport::TestCase
   end
   
   # Diagram ==================================================================
+  test "domain sould return given domain" do
+    domain = Object.new
+    assert_same domain, Class.new(Diagram).new(domain).domain
+  end
+
+  # Diagram DSL ==============================================================
+  test "create should succeed silently if called on abstract class" do
+    create_simple_domain
+    assert_nothing_raised do
+      Diagram.create
+    end
+  end
+
+  test "create should succeed if called on subclass" do
+    create_simple_domain
+    assert_nothing_raised do
+      Class.new(Diagram).create
+    end
+  end
+
+  test "create should call callbacks in instance in specific order" do
+    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
+  end
+
   test "create class method should return result of save" do
     create_simple_domain
     subclass = Class.new(Diagram) do
-      def save
+      save do
         "foobar"
       end
     end
@@ -59,38 +117,13 @@ class DiagramTest < ActiveSupport::TestCase
   test "create should return result of save" do
     create_simple_domain
     diagram = Class.new(Diagram) do
-      def save
+      save do
         "foobar"
       end
     end.new(Domain.generate)
     assert_equal "foobar", diagram.create
   end
   
-  test "domain sould return given domain" do
-    domain = Object.new
-    assert_same domain, Class.new(Diagram).new(domain).domain
-  end
-
-  # Diagram abstractness =====================================================
-  test "create should succeed silently if called on abstract class" do
-    create_simple_domain
-    assert_nothing_raised do
-      Diagram.create
-    end
-  end
-
-  test "create should succeed if called on class that implements process_entity and process_relationship" do
-    create_simple_domain
-    assert_nothing_raised do
-      Class.new(Diagram) do
-        def process_entity(*args)
-        end
-        def process_relationship(*args)
-        end
-      end.create
-    end
-  end
-  
   # Entity filtering =========================================================
   test "generate should yield entities" do
     create_model "Foo"
@@ -111,13 +144,19 @@ class DiagramTest < ActiveSupport::TestCase
     assert_equal [Foo], retrieve_entities(:disconnected => true).map(&:model)
   end
 
-  test "generate should filter descendant entities" do
+  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 descended entities with distinct tables" do
+  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
@@ -126,6 +165,22 @@ class DiagramTest < ActiveSupport::TestCase
     create_table "special_foo", {}, true
     assert_equal [Foo, SpecialFoo], retrieve_entities.map(&:model)
   end
+
+  test "generate should filter generalized entities" do
+    create_model "Cannon"
+    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
@@ -163,27 +218,53 @@ class DiagramTest < ActiveSupport::TestCase
     assert_equal [false, false], retrieve_relationships(:indirect => false).map(&:indirect?)
   end
 
-  test "generate should filter relationships from descendant entities" do
+  test "generate should yield relationships from specialized entities" do
     create_model "Foo", :bar => :references
     create_model "Bar", :type => :string
     Object.const_set :SpecialBar, Class.new(Bar)
     SpecialBar.class_eval do
       has_many :foos
     end
-    assert_equal [], retrieve_relationships
+    assert_equal 1, retrieve_relationships.length
   end
   
-  test "generate should filter relationships to descendant entities" do
+  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
     end
-    assert_equal [], retrieve_relationships
+    assert_equal 1, retrieve_relationships.length
+  end
+
+  # Specialization filtering =================================================
+  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
+
+  test "generate should yield generalizations but not specializations if polymorphism is true" do
+    create_specialization
+    create_generalization
+    assert_equal ["Galleon"], retrieve_specializations(:polymorphism => true).map { |s| s.specialized.name }
+  end
+
+  test "generate should yield specializations and generalizations if polymorphism and inheritance is true" do
+    create_specialization
+    create_generalization
+    assert_equal ["Beer", "Galleon"], retrieve_specializations(:inheritance => true,
+      :polymorphism => true).map { |s| s.specialized.name }
   end
 
   # Attribute filtering ======================================================
-  test "generate should yield regular attributes by default" do
+  test "generate should yield content attributes by default" do
     create_model "Book", :title => :string, :created_at => :datetime, :author => :references do
       belongs_to :author
     end
@@ -222,6 +303,12 @@ class DiagramTest < ActiveSupport::TestCase
     end
     create_model "Author"
     assert_equal %w{created_at title},
-      retrieve_attribute_lists(:attributes => [:regular, :timestamps])[Book].map(&:name)
+      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
 end