diff --git a/tests/test_utils.py b/tests/test_utils.py
index 0597bbf7f8f3e4e1ca23331631a2d3a21e53e77f..6347bd0f5048350c099ba2568dac7caba74baf2d 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1,15 +1,16 @@
 """Test Utils."""
-from typing import List
+from typing import List, Tuple
 
 from attr import attrs, attrib
 
+from flatland.core.grid.grid4 import Grid4TransitionsEnum
 from flatland.envs.rail_env import RailEnvActions
 
 
 @attrs
 class Replay(object):
-    position = attrib()
-    direction = attrib()
+    position = attrib(type=Tuple[int, int])
+    direction = attrib(type=Grid4TransitionsEnum)
     action = attrib(type=RailEnvActions)
     malfunction = attrib(default=0, type=int)
 
@@ -17,5 +18,5 @@ class Replay(object):
 @attrs
 class ReplayConfig(object):
     replay = attrib(type=List[Replay])
-    target = attrib()
+    target = attrib(type=Tuple[int, int])
     speed = attrib(type=float)