From 082e3b7307970fb907bf4d3cfab00305512a7ae2 Mon Sep 17 00:00:00 2001
From: Marek Wolan <marek.wolan@methods.co.uk>
Date: Tue, 30 May 2023 15:48:11 +0100
Subject: [PATCH] Add test for new multidiscrete spaces

---
 .../config/box_obs_space_laydown_config.yaml  | 31 +++++++++++++++++++
 ...ultidiscrete_obs_space_laydown_config.yaml | 31 +++++++++++++++++++
 tests/test_observation_space.py               | 23 ++++++++++++++
 3 files changed, 85 insertions(+)
 create mode 100644 tests/config/box_obs_space_laydown_config.yaml
 create mode 100644 tests/config/multidiscrete_obs_space_laydown_config.yaml
 create mode 100644 tests/test_observation_space.py

diff --git a/tests/config/box_obs_space_laydown_config.yaml b/tests/config/box_obs_space_laydown_config.yaml
new file mode 100644
index 00000000..92226863
--- /dev/null
+++ b/tests/config/box_obs_space_laydown_config.yaml
@@ -0,0 +1,31 @@
+- itemType: ACTIONS
+  type: NODE
+- itemType: OBSERVATIONS
+  type: BOX
+- itemType: STEPS
+  steps: 5
+- itemType: PORTS
+  portsList:
+  - port: '21'
+- itemType: SERVICES
+  serviceList:
+  - name: ftp
+- itemType: NODE
+  node_id: '1'
+  name: node
+  node_class: SERVICE
+  node_type: COMPUTER
+  priority: P1
+  hardware_state: 'ON'
+  ip_address: 192.168.0.1
+  software_state: GOOD
+  file_system_state: GOOD
+  services:
+  - name: ftp
+    port: '21'
+    state: GOOD
+- itemType: POSITION
+  positions:
+  - node: '1'
+    x_pos: 309
+    y_pos: 78
diff --git a/tests/config/multidiscrete_obs_space_laydown_config.yaml b/tests/config/multidiscrete_obs_space_laydown_config.yaml
new file mode 100644
index 00000000..dba1c6de
--- /dev/null
+++ b/tests/config/multidiscrete_obs_space_laydown_config.yaml
@@ -0,0 +1,31 @@
+- itemType: ACTIONS
+  type: NODE
+- itemType: OBSERVATIONS
+  type: MULTIDISCRETE
+- itemType: STEPS
+  steps: 5
+- itemType: PORTS
+  portsList:
+  - port: '21'
+- itemType: SERVICES
+  serviceList:
+  - name: ftp
+- itemType: NODE
+  node_id: '1'
+  name: node
+  node_class: SERVICE
+  node_type: COMPUTER
+  priority: P1
+  hardware_state: 'ON'
+  ip_address: 192.168.0.1
+  software_state: GOOD
+  file_system_state: GOOD
+  services:
+  - name: ftp
+    port: '21'
+    state: GOOD
+- itemType: POSITION
+  positions:
+  - node: '1'
+    x_pos: 309
+    y_pos: 78
diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py
new file mode 100644
index 00000000..2f8e6a24
--- /dev/null
+++ b/tests/test_observation_space.py
@@ -0,0 +1,23 @@
+"""Test env creation and behaviour with different observation spaces."""
+
+from tests import TEST_CONFIG_ROOT
+from tests.conftest import _get_primaite_env_from_config
+
+
+def test_creating_env_with_box_obs():
+    """Try creating env with box observation space."""
+    env = _get_primaite_env_from_config(
+        main_config_path=TEST_CONFIG_ROOT / "one_node_states_on_off_main_config.yaml",
+        lay_down_config_path=TEST_CONFIG_ROOT / "box_obs_space_laydown_config.yaml",
+    )
+    env.update_environent_obs()
+
+
+def test_creating_env_with_multidiscrete_obs():
+    """Try creating env with MultiDiscrete observation space."""
+    env = _get_primaite_env_from_config(
+        main_config_path=TEST_CONFIG_ROOT / "one_node_states_on_off_main_config.yaml",
+        lay_down_config_path=TEST_CONFIG_ROOT
+        / "multidiscrete_obs_space_laydown_config.yaml",
+    )
+    env.update_environent_obs()