diff --git a/tests/test_flatland_envs_observations.py b/tests/test_flatland_envs_observations.py
index 5ee5b4a65afe5dda5a98f6f106e69394a06c71e0..fb45e5fa162d1323533a62349922b569a415395f 100644
--- a/tests/test_flatland_envs_observations.py
+++ b/tests/test_flatland_envs_observations.py
@@ -6,7 +6,7 @@ import numpy as np
 from flatland.envs.generators import rail_from_GridTransitionMap_generator
 from flatland.envs.observations import GlobalObsForRailEnv
 from flatland.envs.rail_env import RailEnv
-from simple_rail import make_simple_rail
+from tests.simple_rail import make_simple_rail
 
 """Tests for `flatland` package."""
 
diff --git a/tests/test_flatland_envs_predictions.py b/tests/test_flatland_envs_predictions.py
index 2926d21d20710c3897cf046550f0977e86b14aa3..96a4b9589093d8635006dac8bea8e902d449600a 100644
--- a/tests/test_flatland_envs_predictions.py
+++ b/tests/test_flatland_envs_predictions.py
@@ -10,7 +10,7 @@ from flatland.envs.observations import TreeObsForRailEnv
 from flatland.envs.predictions import DummyPredictorForRailEnv, ShortestPathPredictorForRailEnv
 from flatland.envs.rail_env import RailEnv
 from flatland.utils.rendertools import RenderTool
-from simple_rail import make_simple_rail
+from tests.simple_rail import make_simple_rail
 
 """Test predictions for `flatland` package."""