diff --git a/tests/test_environments.py b/tests/test_environments.py
index b62a2e60d03b16a89e026a2ac99ea7f9504d6ec3..4f31fe6196846ff843aee3a445d1347d8975799d 100644
--- a/tests/test_environments.py
+++ b/tests/test_environments.py
@@ -1,10 +1,10 @@
 #!/usr/bin/env python
 # -*- coding: utf-8 -*-
+import numpy as np
 
-from flatland.core.env import RailEnv, rail_from_GridTransitionMap_generator
+from flatland.envs.rail_env import RailEnv, rail_from_GridTransitionMap_generator
 from flatland.core.transitions import Grid4Transitions
 from flatland.core.transition_map import GridTransitionMap
-import numpy as np
 
 """Tests for `flatland` package."""