#917 - Integrated the PrimaiteSession into all tests.

- Ran a full pre-commit hook and thus encountered tons of fixes required

tests/test_reward.py

@@ -1,26 +1,33 @@
+import pytest
+
 from tests import TEST_CONFIG_ROOT
-from tests.conftest import _get_primaite_env_from_config
 
 
-def test_rewards_are_being_penalised_at_each_step_function():
+@pytest.mark.parametrize(
+    "temp_primaite_session",
+    [
+        [
+            TEST_CONFIG_ROOT / "one_node_states_on_off_main_config.yaml",
+            TEST_CONFIG_ROOT / "one_node_states_on_off_lay_down_config.yaml",
+        ]
+    ],
+    indirect=True,
+)
+def test_rewards_are_being_penalised_at_each_step_function(
+    temp_primaite_session,
+):
     """
     Test that hardware state is penalised at each step.
 
     When the initial state is OFF compared to reference state which is ON.
-    """
-    env = _get_primaite_env_from_config(
-        training_config_path=TEST_CONFIG_ROOT
-        / "one_node_states_on_off_main_config.yaml",
-        lay_down_config_path=TEST_CONFIG_ROOT
-        / "one_node_states_on_off_lay_down_config.yaml",
-    )
 
-    """
-    On different steps (of the 13 in total) these are the following rewards for config_6 which are activated:
+    On different steps (of the 13 in total) these are the following rewards
+    for config_6 which are activated:
         File System State: goodShouldBeCorrupt = 5 (between Steps 1 & 3)
         Hardware State: onShouldBeOff = -2 (between Steps 4 & 6)
         Service State: goodShouldBeCompromised = 5 (between Steps 7 & 9)
-        Software State (Software State): goodShouldBeCompromised = 5 (between Steps 10 & 12)
+        Software State (Software State): goodShouldBeCompromised = 5 (between
+        Steps 10 & 12)
 
     Total Reward: -2 - 2 + 5 + 5 + 5 + 5 + 5 + 5 = 26
     Step Count: 13
@@ -28,5 +35,8 @@ def test_rewards_are_being_penalised_at_each_step_function():
     For the 4 steps where this occurs the average reward is:
         Average Reward: 2 (26 / 13)
     """
-    print("average reward", env.average_reward)
-    assert env.average_reward == -8.0
+    with temp_primaite_session as session:
+        session.evaluate()
+        session.close()
+        ev_rewards = session.eval_av_reward_per_episode_csv()
+        assert ev_rewards[1] == -8.0