#62 increase unit test coverage

tests/make_simple_rail.py

+import numpy as np
+
+from flatland.core.grid.grid4 import Grid4Transitions
+from flatland.core.transition_map import GridTransitionMap
+
+
+def make_simple_rail():
+    # We instantiate a very simple rail network on a 7x10 grid:
+    #        |
+    #        |
+    #        |
+    # _ _ _ /_\ _ _  _  _ _ _
+    #               \ /
+    #                |
+    #                |
+    #                |
+    cells = [int('0000000000000000', 2),  # empty cell - Case 0
+             int('1000000000100000', 2),  # Case 1 - straight
+             int('1001001000100000', 2),  # Case 2 - simple switch
+             int('1000010000100001', 2),  # Case 3 - diamond drossing
+             int('1001011000100001', 2),  # Case 4 - single slip switch
+             int('1100110000110011', 2),  # Case 5 - double slip switch
+             int('0101001000000010', 2),  # Case 6 - symmetrical switch
+             int('0010000000000000', 2)]  # Case 7 - dead end
+    transitions = Grid4Transitions([])
+    empty = cells[0]
+    dead_end_from_south = cells[7]
+    dead_end_from_west = transitions.rotate_transition(dead_end_from_south, 90)
+    dead_end_from_north = transitions.rotate_transition(dead_end_from_south, 180)
+    dead_end_from_east = transitions.rotate_transition(dead_end_from_south, 270)
+    vertical_straight = cells[1]
+    horizontal_straight = transitions.rotate_transition(vertical_straight, 90)
+    double_switch_south_horizontal_straight = horizontal_straight + cells[6]
+    double_switch_north_horizontal_straight = transitions.rotate_transition(
+        double_switch_south_horizontal_straight, 180)
+    rail_map = np.array(
+        [[empty] * 3 + [dead_end_from_south] + [empty] * 6] +
+        [[empty] * 3 + [vertical_straight] + [empty] * 6] * 2 +
+        [[dead_end_from_east] + [horizontal_straight] * 2 +
+         [double_switch_north_horizontal_straight] +
+         [horizontal_straight] * 2 + [double_switch_south_horizontal_straight] +
+         [horizontal_straight] * 2 + [dead_end_from_west]] +
+        [[empty] * 6 + [vertical_straight] + [empty] * 3] * 2 +
+        [[empty] * 6 + [dead_end_from_north] + [empty] * 3], dtype=np.uint16)
+    rail = GridTransitionMap(width=rail_map.shape[1],
+                             height=rail_map.shape[0], transitions=transitions)
+    rail.grid = rail_map
+    return rail, rail_map

tests/test_flatland_envs_observations.py

@@ -3,62 +3,55 @@
 
 import numpy as np
 
-from flatland.core.transition_map import GridTransitionMap, Grid4Transitions
+from flatland.core.grid.grid4 import Grid4TransitionsEnum
 from flatland.envs.generators import rail_from_GridTransitionMap_generator
-from flatland.envs.observations import GlobalObsForRailEnv
+from flatland.envs.observations import GlobalObsForRailEnv, TreeObsForRailEnv
+from flatland.envs.predictions import DummyPredictorForRailEnv
 from flatland.envs.rail_env import RailEnv
+from flatland.utils.rendertools import RenderTool
+from make_simple_rail import make_simple_rail
 
 """Tests for `flatland` package."""
 
 
+def test_tree_obs(rendering=True):
+    rail, rail_map = make_simple_rail()
+    env = RailEnv(width=rail_map.shape[1],
+                  height=rail_map.shape[0],
+                  rail_generator=rail_from_GridTransitionMap_generator(rail),
+                  number_of_agents=4,
+                  obs_builder_object=TreeObsForRailEnv(max_depth=2, predictor=DummyPredictorForRailEnv(max_depth=10)))
+
+    env.reset()
+
+    # set initial position and direction for testing...
+    env.agents[0].position = (3, 0)
+    env.agents[0].direction = Grid4TransitionsEnum.EAST
+    env.agents[0].target = (3, 9)
+
+    # opposing, conflict
+    env.agents[1].position = (3, 1)
+    env.agents[1].direction = Grid4TransitionsEnum.EAST
+    env.agents[1].target = (3, 0)
+
+    # opposing, conflict???
+    env.agents[2].position = (3, 7)
+    env.agents[2].direction = Grid4TransitionsEnum.EAST
+    env.agents[2].target = (3, 0)
+
+    # same direction
+    env.agents[3].position = (3, 1)
+    env.agents[3].direction = Grid4TransitionsEnum.EAST
+    env.agents[3].target = (3, 9)
+
+    if rendering:
+        renderer = RenderTool(env, gl="PILSVG")
+        renderer.renderEnv(show=True, show_observations=False)
+        input("Continue?")
+
+
 def test_global_obs():
-    # We instantiate a very simple rail network on a 7x10 grid:
-    #        |
-    #        |
-    #        |
-    # _ _ _ /_\ _ _  _  _ _ _
-    #               \ /
-    #                |
-    #                |
-    #                |
-
-    cells = [int('0000000000000000', 2),  # empty cell - Case 0
-             int('1000000000100000', 2),  # Case 1 - straight
-             int('1001001000100000', 2),  # Case 2 - simple switch
-             int('1000010000100001', 2),  # Case 3 - diamond drossing
-             int('1001011000100001', 2),  # Case 4 - single slip switch
-             int('1100110000110011', 2),  # Case 5 - double slip switch
-             int('0101001000000010', 2),  # Case 6 - symmetrical switch
-             int('0010000000000000', 2)]  # Case 7 - dead end
-
-    transitions = Grid4Transitions([])
-    empty = cells[0]
-
-    dead_end_from_south = cells[7]
-    dead_end_from_west = transitions.rotate_transition(dead_end_from_south, 90)
-    dead_end_from_north = transitions.rotate_transition(dead_end_from_south, 180)
-    dead_end_from_east = transitions.rotate_transition(dead_end_from_south, 270)
-
-    vertical_straight = cells[1]
-    horizontal_straight = transitions.rotate_transition(vertical_straight, 90)
-
-    double_switch_south_horizontal_straight = horizontal_straight + cells[6]
-    double_switch_north_horizontal_straight = transitions.rotate_transition(
-        double_switch_south_horizontal_straight, 180)
-
-    rail_map = np.array(
-        [[empty] * 3 + [dead_end_from_south] + [empty] * 6] +
-        [[empty] * 3 + [vertical_straight] + [empty] * 6] * 2 +
-        [[dead_end_from_east] + [horizontal_straight] * 2 +
-         [double_switch_north_horizontal_straight] +
-         [horizontal_straight] * 2 + [double_switch_south_horizontal_straight] +
-         [horizontal_straight] * 2 + [dead_end_from_west]] +
-        [[empty] * 6 + [vertical_straight] + [empty] * 3] * 2 +
-        [[empty] * 6 + [dead_end_from_north] + [empty] * 3], dtype=np.uint16)
-
-    rail = GridTransitionMap(width=rail_map.shape[1],
-                             height=rail_map.shape[0], transitions=transitions)
-    rail.grid = rail_map
+    rail, rail_map = make_simple_rail()
     env = RailEnv(width=rail_map.shape[1],
                   height=rail_map.shape[0],
                   rail_generator=rail_from_GridTransitionMap_generator(rail),

tests/test_flatland_envs_predictions.py

@@ -4,61 +4,17 @@
 import numpy as np
 
 from flatland.core.grid.grid4 import Grid4TransitionsEnum
-from flatland.core.transition_map import GridTransitionMap, Grid4Transitions
 from flatland.envs.generators import rail_from_GridTransitionMap_generator
 from flatland.envs.observations import TreeObsForRailEnv
 from flatland.envs.predictions import DummyPredictorForRailEnv, ShortestPathPredictorForRailEnv
 from flatland.envs.rail_env import RailEnv
 from flatland.envs.rail_env import RailEnvActions
 from flatland.utils.rendertools import RenderTool
+from make_simple_rail import make_simple_rail
 
 """Test predictions for `flatland` package."""
 
 
-def make_simple_rail():
-    # We instantiate a very simple rail network on a 7x10 grid:
-    #        |
-    #        |
-    #        |
-    # _ _ _ /_\ _ _  _  _ _ _
-    #               \ /
-    #                |
-    #                |
-    #                |
-    cells = [int('0000000000000000', 2),  # empty cell - Case 0
-             int('1000000000100000', 2),  # Case 1 - straight
-             int('1001001000100000', 2),  # Case 2 - simple switch
-             int('1000010000100001', 2),  # Case 3 - diamond drossing
-             int('1001011000100001', 2),  # Case 4 - single slip switch
-             int('1100110000110011', 2),  # Case 5 - double slip switch
-             int('0101001000000010', 2),  # Case 6 - symmetrical switch
-             int('0010000000000000', 2)]  # Case 7 - dead end
-    transitions = Grid4Transitions([])
-    empty = cells[0]
-    dead_end_from_south = cells[7]
-    dead_end_from_west = transitions.rotate_transition(dead_end_from_south, 90)
-    dead_end_from_north = transitions.rotate_transition(dead_end_from_south, 180)
-    dead_end_from_east = transitions.rotate_transition(dead_end_from_south, 270)
-    vertical_straight = cells[1]
-    horizontal_straight = transitions.rotate_transition(vertical_straight, 90)
-    double_switch_south_horizontal_straight = horizontal_straight + cells[6]
-    double_switch_north_horizontal_straight = transitions.rotate_transition(
-        double_switch_south_horizontal_straight, 180)
-    rail_map = np.array(
-        [[empty] * 3 + [dead_end_from_south] + [empty] * 6] +
-        [[empty] * 3 + [vertical_straight] + [empty] * 6] * 2 +
-        [[dead_end_from_east] + [horizontal_straight] * 2 +
-         [double_switch_north_horizontal_straight] +
-         [horizontal_straight] * 2 + [double_switch_south_horizontal_straight] +
-         [horizontal_straight] * 2 + [dead_end_from_west]] +
-        [[empty] * 6 + [vertical_straight] + [empty] * 3] * 2 +
-        [[empty] * 6 + [dead_end_from_north] + [empty] * 3], dtype=np.uint16)
-    rail = GridTransitionMap(width=rail_map.shape[1],
-                             height=rail_map.shape[0], transitions=transitions)
-    rail.grid = rail_map
-    return rail, rail_map
-
-
 def test_dummy_predictor(rendering=False):
     rail, rail_map = make_simple_rail()