Adding TrafficLightStateMachine and TrafficLightController

tests/resources/traffic_lights/intersection_1.json

@@ -0,0 +1,57 @@
+[
+    {
+        "actor_states": {
+            "4411": "red",
+            "3411": "red",
+            "4399": "red",
+            "3399": "red"
+        },
+        "state": "0",
+        "duration": "10",
+        "next_state": "1"
+    },
+    {
+        "actor_states": {
+            "4411": "green",
+            "3411": "red",
+            "4399": "green",
+            "3399": "green"
+        },
+        "state": "1",
+        "duration": "10",
+        "next_state": "2"
+    },
+    {
+        "actor_states": {
+            "4411": "green",
+            "3411": "red",
+            "4399": "yellow",
+            "3399": "yellow"
+        },
+        "state": "2",
+        "duration": "5",
+        "next_state": "3"
+    },
+    {
+        "actor_states": {
+            "4411": "green",
+            "3411": "green",
+            "4399": "red",
+            "3399": "red"
+        },
+        "state": "3",
+        "duration": "10",
+        "next_state": "4"
+    },
+    {
+        "actor_states": {
+            "4411": "yellow",
+            "3411": "yellow",
+            "4399": "red",
+            "3399": "red"
+        },
+        "state": "4",
+        "duration": "5",
+        "next_state": "0"
+    }
+]

tests/resources/traffic_lights/intersection_2.json

@@ -0,0 +1,93 @@
+[
+    {
+        "actor_states": {
+            "3404": "red",
+            "4404": "red",
+            "3406": "red",
+            "3405": "red",
+            "3403": "red",
+            "4403": "red"
+        },
+        "state": "0",
+        "duration": "2",
+        "next_state": "1"
+    },
+    {
+        "actor_states": {
+            "3404": "red",
+            "4404": "red",
+            "3406": "green",
+            "3405": "green",
+            "3403": "red",
+            "4403": "red"
+        },
+        "state": "1",
+        "duration": "10",
+        "next_state": "2"
+    },
+    {
+        "actor_states": {
+            "3404": "red",
+            "4404": "red",
+            "3406": "yellow",
+            "3405": "yellow",
+            "3403": "red",
+            "4403": "red"
+        },
+        "state": "2",
+        "duration": "5",
+        "next_state": "3"
+    },
+    {
+        "actor_states": {
+            "3404": "green",
+            "4404": "green",
+            "3406": "red",
+            "3405": "red",
+            "3403": "red",
+            "4403": "red"
+        },
+        "state": "3",
+        "duration": "10",
+        "next_state": "4"
+    },
+    {
+        "actor_states": {
+            "3404": "yellow",
+            "4404": "yellow",
+            "3406": "red",
+            "3405": "red",
+            "3403": "red",
+            "4403": "red"
+        },
+        "state": "4",
+        "duration": "5",
+        "next_state": "5"
+    },
+    {
+        "actor_states": {
+            "3404": "red",
+            "4404": "red",
+            "3406": "red",
+            "3405": "red",
+            "3403": "green",
+            "4403": "green"
+        },
+        "state": "5",
+        "duration": "10",
+        "next_state": "6"
+    },
+    {
+        "actor_states": {
+            "3404": "red",
+            "4404": "red",
+            "3406": "red",
+            "3405": "red",
+            "3403": "yellow",
+            "4403": "yellow"
+        },
+        "state": "6",
+        "duration": "5",
+        "next_state": "0"
+    }
+]

tests/test_traffic_light_controller.py

@@ -0,0 +1,75 @@
+import pytest
+import os
+from torchdrivesim.traffic_lights import TrafficLightController, TrafficLightStateMachine
+
+@pytest.fixture
+def traffic_light_controller():
+    folder_path = os.path.join(os.path.dirname(__file__), "resources",  "traffic_lights")
+    state_machine_files = [os.path.join(folder_path, f) for f in os.listdir(folder_path) if f.endswith(".json")]
+    state_machine_files.sort()
+    light_machines = [TrafficLightStateMachine(f) for f in state_machine_files]
+    return TrafficLightController(light_machines)
+
+
+def test_reset(traffic_light_controller: TrafficLightController):
+    traffic_light_controller.reset()
+    assert all([t > 0 for t in traffic_light_controller.time_remaining])
+
+def test_set_to(traffic_light_controller: TrafficLightController):
+    traffic_light_controller.set_to([[2, 4],[0, 1]])
+    assert traffic_light_controller.time_remaining == [4, 1]
+    assert traffic_light_controller.current_state == {
+                "4411": "green",
+                "3411": "red",
+                "4399": "yellow",
+                "3399": "yellow",
+                "3404": "red",
+                "4404": "red",
+                "3406": "red",
+                "3405": "red",
+                "3403": "red",
+                "4403": "red"
+    }
+    traffic_light_controller.set_to([[2, 3],[1, 0]])
+    assert traffic_light_controller.time_remaining == [3, 0]
+    assert traffic_light_controller.current_state == {
+                "4411": "green",
+                "3411": "red",
+                "4399": "yellow",
+                "3399": "yellow",
+                "3404": "red",
+                "4404": "red",
+                "3406": "green",
+                "3405": "green",
+                "3403": "red",
+                "4403": "red"
+    }
+
+def test_tick(traffic_light_controller: TrafficLightController):
+    traffic_light_controller.set_to([[0, 10], [0, 2]])
+    traffic_light_controller.tick(10)
+    assert traffic_light_controller.current_state == {
+        "4411": "red",
+        "3411": "red",
+        "4399": "red",
+        "3399": "red",
+        "3404": "red",
+        "4404": "red",
+        "3406": "red",
+        "3405": "red",
+        "3403": "red",
+        "4403": "red"
+    }
+    traffic_light_controller.tick(10)
+    assert traffic_light_controller.current_state == {
+        "4411": "red",
+        "3411": "red",
+        "4399": "red",
+        "3399": "red",
+        "3404": "red",
+        "4404": "red",
+        "3406": "green",
+        "3405": "green",
+        "3403": "red",
+        "4403": "red"
+    }

tests/test_traffic_light_state_machine.py

@@ -0,0 +1,80 @@
+import os
+import pytest
+from torchdrivesim.traffic_lights import TrafficLightStateMachine
+
+@pytest.fixture
+def traffic_light_state_machine():
+    state_machine_file_path = os.path.join(os.path.dirname(__file__), "resources", "traffic_lights", 'intersection_1.json')
+    return TrafficLightStateMachine(state_machine_file_path)
+
+def test_reset(traffic_light_state_machine: TrafficLightStateMachine):
+    traffic_light_state_machine.reset()
+    assert traffic_light_state_machine.time_remaining >= 1
+
+def test_set_to(traffic_light_state_machine: TrafficLightStateMachine):
+    traffic_light_state_machine.set_to(2, time_remaining=3)
+    assert traffic_light_state_machine.time_remaining == 3
+    assert traffic_light_state_machine.current_state ==  {
+            "actor_states": {
+                "4411": "green",
+                "3411": "red",
+                "4399": "yellow",
+                "3399": "yellow"
+            },
+            "state": "2",
+            "duration": 5,
+            "next_state": "3"
+        }
+    traffic_light_state_machine.set_to(2, time_remaining=1)
+    assert traffic_light_state_machine.time_remaining == 1
+    assert traffic_light_state_machine.current_state == {
+            "actor_states": {
+                "4411": "green",
+                "3411": "red",
+                "4399": "yellow",
+                "3399": "yellow"
+            },
+            "state": "2",
+            "duration": 5,
+            "next_state": "3"
+        }
+
+
+def test_tick(traffic_light_state_machine: TrafficLightStateMachine):
+    traffic_light_state_machine.set_to(0, 1)
+    traffic_light_state_machine.tick(9)
+    assert traffic_light_state_machine.time_remaining <= 0.1
+    assert traffic_light_state_machine.current_state == {
+            "actor_states": {
+                "4411": "red",
+                "3411": "red",
+                "4399": "red",
+                "3399": "red"
+            },
+            "state": "0",
+            "duration": 10,
+            "next_state": "1"
+        }
+    traffic_light_state_machine.tick(1)
+    assert traffic_light_state_machine.time_remaining == 10
+    assert traffic_light_state_machine.current_state == {
+            "actor_states": {
+                "4411": "green",
+                "3411": "red",
+                "4399": "green",
+                "3399": "green"
+            },
+            "state": "1",
+            "duration": 10,
+            "next_state": "2"
+        }
+
+def test_get_current_actor_states(traffic_light_state_machine: TrafficLightStateMachine):
+    traffic_light_state_machine.set_to(4, 3)
+    actor_states = traffic_light_state_machine.get_current_actor_states()
+    assert actor_states == {
+                "4411": "yellow",
+                "3411": "yellow",
+                "4399": "red",
+                "3399": "red"
+            }