gh-1612 - Tests added for new functionality

gchq · Jan 16, 2018 · 9906111 · 9906111
1 parent 0b7d5fd
commit 9906111
Showing 1 changed file with 43 additions and 0 deletions.
diff --git a/core/data/src/test/java/uk/gov/gchq/gaffer/data/graph/WalkTest.java b/core/data/src/test/java/uk/gov/gchq/gaffer/data/graph/WalkTest.java
@@ -44,6 +44,7 @@
 import static org.hamcrest.core.Is.is;
 import static org.hamcrest.core.IsCollectionContaining.hasItems;
 import static org.hamcrest.core.IsInstanceOf.instanceOf;
+import static org.junit.Assert.assertEquals;
 import static org.junit.Assert.fail;
 
 public class WalkTest {
@@ -501,6 +502,48 @@ public void shouldGetNotPath() {
         assertThat(walk.isPath(), is(false));
     }
 
+    @Test
+    public void shouldGetSourceVertexFromWalk() {
+        // Given
+        // [A] -> [B] -> [C]
+        //  \             \
+        //   (BasicEntity) (BasicEntity)
+
+        final Walk walk = new Walk.Builder()
+                .entity(ENTITY_A)
+                .edges(EDGE_AB, EDGE_BC)
+                .entity(ENTITY_C)
+                .build();
+
+        // When
+        final Object result = walk.getSourceVertex();
+
+        // Then
+        assertEquals(result, "A");
+    }
+
+    @Test
+    public void shouldGetDestinationVertexFromWalk() {
+        // Given
+        // [A]     ->    [E]     ->    [D]
+        //  \             \             \
+        //   (BasicEntity) (BasicEntity) (BasicEntity)
+
+        final Walk walk = new Walk.Builder()
+                .entity(ENTITY_A)
+                .edge(EDGE_AE)
+                .entities(ENTITY_E)
+                .edge(EDGE_ED)
+                .entity(ENTITY_D)
+                .build();
+
+        // When
+        final Object result = walk.getDestinationVertex();
+
+        // Then
+        assertEquals(result, "D");
+    }
+
     private List<EdgeId> getEdges() {
         return IntStream.range(0, 10).mapToObj(i -> {
             return new Edge.Builder()