Merge branch 'dev' into feature/2041_2042-Add-NTP-Services

tests/conftest.py

@@ -1,17 +1,13 @@
 # © Crown-owned copyright 2023, Defence Science and Technology Laboratory UK
-import datetime
-import shutil
-import tempfile
-from datetime import datetime
 from pathlib import Path
 from typing import Any, Dict, Union
 
-import nodeenv
 import pytest
 import yaml
 
 from primaite import getLogger
-from primaite.game.session import PrimaiteSession
+from primaite.game.game import PrimaiteGame
+from primaite.session.session import PrimaiteSession
 
 # from primaite.environment.primaite_env import Primaite
 # from primaite.primaite_session import PrimaiteSession

tests/e2e_integration_tests/environments/test_rllib_multi_agent_environment.py

@@ -0,0 +1,45 @@
+import pytest
+import ray
+import yaml
+from ray import air, tune
+from ray.rllib.algorithms.ppo import PPOConfig
+
+from primaite.config.load import example_config_path
+from primaite.game.game import PrimaiteGame
+from primaite.session.environment import PrimaiteRayMARLEnv
+
+
+@pytest.mark.skip(reason="Slow, reenable later")
+def test_rllib_multi_agent_compatibility():
+    """Test that the PrimaiteRayEnv class can be used with a multi agent RLLIB system."""
+
+    with open(example_config_path(), "r") as f:
+        cfg = yaml.safe_load(f)
+
+    game = PrimaiteGame.from_config(cfg)
+
+    ray.shutdown()
+    ray.init()
+
+    env_config = {"game": game}
+    config = (
+        PPOConfig()
+        .environment(env=PrimaiteRayMARLEnv, env_config={"game": game})
+        .rollouts(num_rollout_workers=0)
+        .multi_agent(
+            policies={agent.agent_name for agent in game.rl_agents},
+            policy_mapping_fn=lambda agent_id, episode, worker, **kw: agent_id,
+        )
+        .training(train_batch_size=128)
+    )
+
+    tune.Tuner(
+        "PPO",
+        run_config=air.RunConfig(
+            stop={"training_iteration": 128},
+            checkpoint_config=air.CheckpointConfig(
+                checkpoint_frequency=10,
+            ),
+        ),
+        param_space=config,
+    ).fit()

tests/e2e_integration_tests/environments/test_rllib_single_agent_environment.py

@@ -0,0 +1,40 @@
+import tempfile
+from pathlib import Path
+
+import pytest
+import ray
+import yaml
+from ray.rllib.algorithms import ppo
+
+from primaite.config.load import example_config_path
+from primaite.game.game import PrimaiteGame
+from primaite.session.environment import PrimaiteRayEnv
+
+
+@pytest.mark.skip(reason="Slow, reenable later")
+def test_rllib_single_agent_compatibility():
+    """Test that the PrimaiteRayEnv class can be used with a single agent RLLIB system."""
+    with open(example_config_path(), "r") as f:
+        cfg = yaml.safe_load(f)
+
+    game = PrimaiteGame.from_config(cfg)
+
+    ray.shutdown()
+    ray.init()
+
+    env_config = {"game": game}
+    config = {
+        "env": PrimaiteRayEnv,
+        "env_config": env_config,
+        "disable_env_checking": True,
+        "num_rollout_workers": 0,
+    }
+
+    algo = ppo.PPO(config=config)
+
+    for i in range(5):
+        result = algo.train()
+
+    save_file = Path(tempfile.gettempdir()) / "ray/"
+    algo.save(save_file)
+    assert save_file.exists()

tests/e2e_integration_tests/environments/test_sb3_environment.py

@@ -0,0 +1,27 @@
+"""Test that we can create a primaite environment and train sb3 agent with no crash."""
+import tempfile
+from pathlib import Path
+
+import yaml
+from stable_baselines3 import PPO
+
+from primaite.config.load import example_config_path
+from primaite.game.game import PrimaiteGame
+from primaite.session.environment import PrimaiteGymEnv
+
+
+def test_sb3_compatibility():
+    """Test that the Gymnasium environment can be used with an SB3 agent."""
+    with open(example_config_path(), "r") as f:
+        cfg = yaml.safe_load(f)
+
+    game = PrimaiteGame.from_config(cfg)
+    gym = PrimaiteGymEnv(game=game)
+    model = PPO("MlpPolicy", gym)
+
+    model.learn(total_timesteps=1000)
+
+    save_path = Path(tempfile.gettempdir()) / "model.zip"
+    model.save(save_path)
+
+    assert (save_path).exists()

tests/e2e_integration_tests/test_primaite_session.py

@@ -1,12 +1,14 @@
 import pydantic
 import pytest
 
+from tests import TEST_ASSETS_ROOT
 from tests.conftest import TempPrimaiteSession
 
-CFG_PATH = "tests/assets/configs/test_primaite_session.yaml"
-TRAINING_ONLY_PATH = "tests/assets/configs/train_only_primaite_session.yaml"
-EVAL_ONLY_PATH = "tests/assets/configs/eval_only_primaite_session.yaml"
-MISCONFIGURED_PATH = "tests/assets/configs/bad_primaite_session.yaml"
+CFG_PATH = TEST_ASSETS_ROOT / "configs/test_primaite_session.yaml"
+TRAINING_ONLY_PATH = TEST_ASSETS_ROOT / "configs/train_only_primaite_session.yaml"
+EVAL_ONLY_PATH = TEST_ASSETS_ROOT / "configs/eval_only_primaite_session.yaml"
+MISCONFIGURED_PATH = TEST_ASSETS_ROOT / "configs/bad_primaite_session.yaml"
+MULTI_AGENT_PATH = TEST_ASSETS_ROOT / "configs/multi_agent_session.yaml"
 
 
 class TestPrimaiteSession:
@@ -18,15 +20,15 @@ class TestPrimaiteSession:
                 raise AssertionError
 
             assert session is not None
-            assert session.simulation
-            assert len(session.agents) == 3
-            assert len(session.rl_agents) == 1
+            assert session.game.simulation
+            assert len(session.game.agents) == 3
+            assert len(session.game.rl_agents) == 1
 
             assert session.policy
             assert session.env
 
-            assert session.simulation.network
-            assert len(session.simulation.network.nodes) == 10
+            assert session.game.simulation.network
+            assert len(session.game.simulation.network.nodes) == 10
 
     @pytest.mark.parametrize("temp_primaite_session", [[CFG_PATH]], indirect=True)
     def test_start_session(self, temp_primaite_session):
@@ -63,6 +65,27 @@ class TestPrimaiteSession:
             session.start_session()
             # TODO: include checks that the model was loaded and that the eval-only session ran
 
+    @pytest.mark.skip(reason="Slow, reenable later")
+    @pytest.mark.parametrize("temp_primaite_session", [[MULTI_AGENT_PATH]], indirect=True)
+    def test_multi_agent_session(self, temp_primaite_session):
+        """Check that we can run a training session with a multi agent system."""
+        with temp_primaite_session as session:
+            session.start_session()
+
     def test_error_thrown_on_bad_configuration(self):
         with pytest.raises(pydantic.ValidationError):
             session = TempPrimaiteSession.from_config(MISCONFIGURED_PATH)
+
+    @pytest.mark.parametrize("temp_primaite_session", [[CFG_PATH]], indirect=True)
+    def test_session_sim_reset(self, temp_primaite_session):
+        with temp_primaite_session as session:
+            session: TempPrimaiteSession
+            client_1 = session.game.simulation.network.get_node_by_hostname("client_1")
+            client_1.software_manager.uninstall("DataManipulationBot")
+
+            assert "DataManipulationBot" not in client_1.software_manager.software
+
+            session.game.reset()
+            client_1 = session.game.simulation.network.get_node_by_hostname("client_1")
+
+            assert "DataManipulationBot" in client_1.software_manager.software

tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py

@@ -19,16 +19,16 @@ def test_data_manipulation(uc2_network):
     db_service.backup_database()
 
     # First check that the DB client on the web_server can successfully query the users table on the database
-    assert db_client.query("SELECT * FROM user;")
+    assert db_client.query("SELECT")
 
     # Now we run the DataManipulationBot
     db_manipulation_bot.run()
 
     # Now check that the DB client on the web_server cannot query the users table on the database
-    assert not db_client.query("SELECT * FROM user;")
+    assert not db_client.query("SELECT")
 
     # Now restore the database
     db_service.restore_backup()
 
     # Now check that the DB client on the web_server can successfully query the users table on the database
-    assert db_client.query("SELECT * FROM user;")
+    assert db_client.query("SELECT")

tests/integration_tests/system/test_database_on_node.py

@@ -57,7 +57,7 @@ def test_database_client_query(uc2_network):
     db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
     db_client.connect()
 
-    assert db_client.query("SELECT * FROM user;")
+    assert db_client.query("SELECT")
 
 
 def test_create_database_backup(uc2_network):

tests/unit_tests/_primaite/_simulator/_system/_services/_red_services/test_data_manipulation_bot.py

@@ -17,4 +17,4 @@ def test_creation():
     assert data_manipulation_bot.name == "DataManipulationBot"
     assert data_manipulation_bot.port == Port.POSTGRES_SERVER
     assert data_manipulation_bot.protocol == IPProtocol.TCP
-    assert data_manipulation_bot.payload == "DROP TABLE IF EXISTS user;"
+    assert data_manipulation_bot.payload == "DELETE"