Add weighted value for edges

lib/mosaik/graph/graph.rb

@@ -17,6 +17,14 @@ def add_vertex(vertex)
         vertices[vertex.value] = vertex
       end
 
+      def add_edge(from, to, weight = nil)
+        vertices[from.value].add_edge(to, weight)
+      end
+
+      def remove_edge(from, to)
+        vertices[from.value].remove_edge(to)
+      end
+
       def find_vertex(value)
         vertices[value]
       end

lib/mosaik/graph/vertex.rb

@@ -5,13 +5,17 @@ module Graph
     class Vertex
       attr_reader :value, :edges
 
-      def initialize(value, edges = Set.new)
+      def initialize(value, edges = {})
         @value = value
         @edges = edges
       end
 
-      def add_edge(vertex)
-        edges << vertex.value
+      def add_edge(vertex, weight = nil)
+        edges[vertex.value] = weight
+      end
+
+      def remove_edge(vertex)
+        edges.delete(vertex.value)
       end
 
       def inspect

spec/mosaik/graph/graph_spec.rb

@@ -13,6 +13,112 @@
     end
   end
 
+  describe "#add_edge" do
+    it "adds an edge" do
+      vertex1 = build(:vertex, value: "vertex1")
+      vertex2 = build(:vertex, value: "vertex2")
+
+      graph.add_vertex(vertex1)
+      graph.add_vertex(vertex2)
+      graph.add_edge(vertex1, vertex2)
+
+      expect(graph.find_vertex("vertex1").edges.keys).to eq ["vertex2"]
+    end
+
+    it "does not add an edge twice" do
+      vertex1 = build(:vertex, value: "vertex1")
+      vertex2 = build(:vertex, value: "vertex2")
+
+      graph.add_vertex(vertex1)
+      graph.add_vertex(vertex2)
+      graph.add_edge(vertex1, vertex2)
+      graph.add_edge(vertex1, vertex2)
+
+      expect(graph.find_vertex("vertex1").edges.keys).to eq ["vertex2"]
+    end
+
+    describe "weighted" do
+      it "adds a weighted edge" do
+        vertex1 = build(:vertex, value: "vertex1")
+        vertex2 = build(:vertex, value: "vertex2")
+
+        graph.add_vertex(vertex1)
+        graph.add_vertex(vertex2)
+        graph.add_edge(vertex1, vertex2, 5)
+
+        expect(graph.find_vertex("vertex1").edges.keys).to eq ["vertex2"]
+        expect(graph.find_vertex("vertex1").edges["vertex2"]).to eq 5
+      end
+
+      it "overrides a weighted edge" do
+        vertex1 = build(:vertex, value: "vertex1")
+        vertex2 = build(:vertex, value: "vertex2")
+
+        graph.add_vertex(vertex1)
+        graph.add_vertex(vertex2)
+        graph.add_edge(vertex1, vertex2, 5)
+        graph.add_edge(vertex1, vertex2, 10)
+
+        expect(graph.find_vertex("vertex1").edges.keys).to eq ["vertex2"]
+        expect(graph.find_vertex("vertex1").edges["vertex2"]).to eq 10
+      end
+    end
+  end
+
+  describe "#remove_edge" do
+    it "removes an edge" do
+      vertex1 = build(:vertex, value: "vertex1")
+      vertex2 = build(:vertex, value: "vertex2")
+
+      graph.add_vertex(vertex1)
+      graph.add_vertex(vertex2)
+      graph.add_edge(vertex1, vertex2)
+
+      graph.remove_edge(vertex1, vertex2)
+
+      expect(graph.find_vertex("vertex1").edges).to be_empty
+    end
+
+    it "does not remove a non-existing edge" do
+      vertex1 = build(:vertex, value: "vertex1")
+      vertex2 = build(:vertex, value: "vertex2")
+
+      graph.add_vertex(vertex1)
+      graph.add_vertex(vertex2)
+
+      graph.remove_edge(vertex1, vertex2)
+
+      expect(graph.find_vertex("vertex1").edges).to be_empty
+    end
+
+    describe "weighted" do
+      it "removes a weighted edge" do
+        vertex1 = build(:vertex, value: "vertex1")
+        vertex2 = build(:vertex, value: "vertex2")
+
+        graph.add_vertex(vertex1)
+        graph.add_vertex(vertex2)
+        graph.add_edge(vertex1, vertex2, 5)
+
+        graph.remove_edge(vertex1, vertex2)
+
+        expect(graph.find_vertex("vertex1").edges).to be_empty
+      end
+
+      it "does not remove a non-existing weighted edge" do
+        vertex1 = build(:vertex, value: "vertex1")
+        vertex2 = build(:vertex, value: "vertex2")
+
+        graph.add_vertex(vertex1)
+        graph.add_vertex(vertex2)
+
+        graph.remove_edge(vertex1, vertex2)
+
+        expect(graph.find_vertex("vertex1").edges).to be_empty
+      end
+    end
+  end
+
   describe "#find_vertex" do
     it "finds a vertex" do
       graph.add_vertex(build(:vertex, value: "address"))

spec/mosaik/graph/vertex_spec.rb

@@ -9,7 +9,7 @@
 
       vertex.add_edge(child)
 
-      expect(vertex.edges.to_a).to eq ["child"]
+      expect(vertex.edges.keys).to eq ["child"]
     end
 
     it "does not add an edge twice" do
@@ -18,7 +18,28 @@
       vertex.add_edge(child)
       vertex.add_edge(child)
 
-      expect(vertex.edges.to_a).to eq ["child"]
+      expect(vertex.edges.keys).to eq ["child"]
+    end
+
+    describe "weighted edge" do
+      it "adds a weighted edge" do
+        child = build(:vertex, value: "child")
+
+        vertex.add_edge(child, 5)
+
+        expect(vertex.edges.keys).to eq ["child"]
+        expect(vertex.edges["child"]).to eq 5
+      end
+
+      it "overrides a weighted edge" do
+        child = build(:vertex, value: "child")
+
+        vertex.add_edge(child, 5)
+        vertex.add_edge(child, 10)
+
+        expect(vertex.edges.keys).to eq ["child"]
+        expect(vertex.edges["child"]).to eq 10
+      end
     end
   end
 end