From 8ac77800a0a6a3f1e114fab18d7a721949302780 Mon Sep 17 00:00:00 2001
From: Loredana <lore.crusoveanu@gmail.com>
Date: Sun, 26 May 2019 19:38:40 +0300
Subject: [PATCH] fix remove vertex

---
 .../main/java/com/baeldung/graph/Graph.java   |  2 +-
 ...versalUnitTest.java => GraphUnitTest.java} | 19 +++++++++++++++----
 2 files changed, 16 insertions(+), 5 deletions(-)
 rename core-java-modules/core-java/src/test/java/com/baeldung/graph/{GraphTraversalUnitTest.java => GraphUnitTest.java} (60%)

diff --git a/core-java-modules/core-java/src/main/java/com/baeldung/graph/Graph.java b/core-java-modules/core-java/src/main/java/com/baeldung/graph/Graph.java
index 3f2e17c43c..58713b1b3f 100644
--- a/core-java-modules/core-java/src/main/java/com/baeldung/graph/Graph.java
+++ b/core-java-modules/core-java/src/main/java/com/baeldung/graph/Graph.java
@@ -19,7 +19,7 @@ public class Graph {
 
     void removeVertex(String label) {
         Vertex v = new Vertex(label);
-        adjVertices.values().stream().map(e -> e.remove(v)).collect(Collectors.toList());
+        adjVertices.values().stream().forEach(e -> e.remove(v));
         adjVertices.remove(new Vertex(label));
     }
 
diff --git a/core-java-modules/core-java/src/test/java/com/baeldung/graph/GraphTraversalUnitTest.java b/core-java-modules/core-java/src/test/java/com/baeldung/graph/GraphUnitTest.java
similarity index 60%
rename from core-java-modules/core-java/src/test/java/com/baeldung/graph/GraphTraversalUnitTest.java
rename to core-java-modules/core-java/src/test/java/com/baeldung/graph/GraphUnitTest.java
index d955d56d95..68611e508b 100644
--- a/core-java-modules/core-java/src/test/java/com/baeldung/graph/GraphTraversalUnitTest.java
+++ b/core-java-modules/core-java/src/test/java/com/baeldung/graph/GraphUnitTest.java
@@ -1,20 +1,31 @@
 package com.baeldung.graph;
 
-import org.junit.Assert;
+import static org.junit.Assert.assertEquals;
 import org.junit.Test;
 
-public class GraphTraversalUnitTest {
+public class GraphUnitTest {
     @Test
     public void givenAGraph_whenTraversingDepthFirst_thenExpectedResult() {
         Graph graph = createGraph();
-        Assert.assertEquals("[Bob, Rob, Maria, Alice, Mark]", 
+        assertEquals("[Bob, Rob, Maria, Alice, Mark]", 
             GraphTraversal.depthFirstTraversal(graph, "Bob").toString());
     }
     
     @Test
     public void givenAGraph_whenTraversingBreadthFirst_thenExpectedResult() {
         Graph graph = createGraph();
-        Assert.assertEquals("[Bob, Alice, Rob, Mark, Maria]", 
+       assertEquals("[Bob, Alice, Rob, Mark, Maria]", 
+            GraphTraversal.breadthFirstTraversal(graph, "Bob").toString());
+    }
+    
+    @Test
+    public void givenAGraph_whenRemoveVertex_thenVertedNotFound() {
+        Graph graph = createGraph();
+        assertEquals("[Bob, Alice, Rob, Mark, Maria]", 
+            GraphTraversal.breadthFirstTraversal(graph, "Bob").toString());
+        
+        graph.removeVertex("Maria");
+        assertEquals("[Bob, Alice, Rob, Mark]", 
             GraphTraversal.breadthFirstTraversal(graph, "Bob").toString());
     }