diff --git a/src/test/java/com/arangodb/graphql/schema/ArangoEdgeDirectiveTest.java b/src/test/java/com/arangodb/graphql/schema/ArangoEdgeDirectiveTest.java
index f58c5cd..e6d6ffe 100644
--- a/src/test/java/com/arangodb/graphql/schema/ArangoEdgeDirectiveTest.java
+++ b/src/test/java/com/arangodb/graphql/schema/ArangoEdgeDirectiveTest.java
@@ -55,6 +55,7 @@ public void testCollection(){
         String expected = "collectionName";
 
         when(field.getDirective(eq("edge"))).thenReturn(schemaDirective);
+        when(schemaDirective.getArgument(eq("direction"))).thenReturn(null);
         when(schemaDirective.getArgument(eq("collection"))).thenReturn(argument);
         when(argument.getValue()).thenReturn(expected);
 
@@ -79,6 +80,7 @@ public void testNoDirective(){
     public void testNoCollection(){
 
         when(field.getDirective(eq("edge"))).thenReturn(schemaDirective);
+        when(schemaDirective.getArgument(eq("direction"))).thenReturn(null);
         when(schemaDirective.getArgument(eq("collection"))).thenReturn(null);
 
         directive = new ArangoEdgeDirective(field);