From 5d0d968af54bf21b892c25074a9ca452604bb210 Mon Sep 17 00:00:00 2001
From: Neil Smith <neil.github@njae.me.uk>
Date: Fri, 23 Sep 2011 22:25:10 +0100
Subject: [PATCH] Handles subclasses of Vertex and Edge

---
 lib/graph.njae/graph.rb  |  4 ++--
 spec/graph/graph_spec.rb | 17 +++++++++++++++++
 2 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/lib/graph.njae/graph.rb b/lib/graph.njae/graph.rb
index 1da5297..e08a1ad 100644
--- a/lib/graph.njae/graph.rb
+++ b/lib/graph.njae/graph.rb
@@ -16,9 +16,9 @@ module GraphNjae
     
     # Add a Vertex or Edge to the graph.
     def <<(other)
-      if other.class == Vertex
+      if other.class.ancestors.include? Vertex
         self.vertices << other
-      elsif
+      elsif other.class.ancestors.include? Edge
         self.edges << other
       end
       self
diff --git a/spec/graph/graph_spec.rb b/spec/graph/graph_spec.rb
index ba7bea1..09cc6a4 100644
--- a/spec/graph/graph_spec.rb
+++ b/spec/graph/graph_spec.rb
@@ -1,6 +1,13 @@
 require File.expand_path(File.dirname(__FILE__) + '/../spec_helper')
 
 module GraphNjae
+
+  class SVertex < Vertex
+  end
+  
+  class SEdge < Edge
+  end
+  
   describe Graph do
     let (:g) { Graph.new }
     describe "#initialize" do
@@ -54,6 +61,16 @@ module GraphNjae
         g.edges.should include(e1)
         g.edges.should include(e2)
       end
+      
+      it "adds a subclass of Vertex" do
+        g.vertices.should be_empty
+        v1 = SVertex.new
+        v2 = SVertex.new
+        g << v1 << v2
+        g.should have(2).vertices
+        g.vertices.should include(v1)
+        g.vertices.should include(v2)        
+      end
     end # #<<
 
     describe "connect" do
-- 
2.34.1