From 5d0d968af54bf21b892c25074a9ca452604bb210 Mon Sep 17 00:00:00 2001 From: Neil Smith 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