From 76ec9683cb96135bd35cc297b5531b1141e0aeb8 Mon Sep 17 00:00:00 2001
From: Marek Wolan <marek.wolan@methods.co.uk>
Date: Thu, 1 Jun 2023 09:45:46 +0100
Subject: [PATCH] Improve observation space test

---
 .../config/box_obs_space_laydown_config.yaml  | 63 +++++++++++++++----
 ...ultidiscrete_obs_space_laydown_config.yaml | 63 +++++++++++++++----
 tests/test_observation_space.py               | 11 ++++
 3 files changed, 111 insertions(+), 26 deletions(-)

diff --git a/tests/config/box_obs_space_laydown_config.yaml b/tests/config/box_obs_space_laydown_config.yaml
index 92226863..203bc0e7 100644
--- a/tests/config/box_obs_space_laydown_config.yaml
+++ b/tests/config/box_obs_space_laydown_config.yaml
@@ -6,26 +6,63 @@
   steps: 5
 - itemType: PORTS
   portsList:
-  - port: '21'
+    - port: '80'
 - itemType: SERVICES
   serviceList:
-  - name: ftp
+    - name: TCP
+
+########################################
+# Nodes
 - itemType: NODE
   node_id: '1'
-  name: node
+  name: PC1
   node_class: SERVICE
   node_type: COMPUTER
-  priority: P1
+  priority: P5
   hardware_state: 'ON'
-  ip_address: 192.168.0.1
+  ip_address: 192.168.1.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
+    - name: TCP
+      port: '80'
+      state: GOOD
+- itemType: NODE
+  node_id: '2'
+  name: SERVER
+  node_class: SERVICE
+  node_type: SERVER
+  priority: P5
+  hardware_state: 'ON'
+  ip_address: 192.168.1.2
+  software_state: GOOD
+  file_system_state: GOOD
+  services:
+    - name: TCP
+      port: '80'
+      state: GOOD
+- itemType: NODE
+  node_id: '3'
+  name: SWITCH1
+  node_class: ACTIVE
+  node_type: SWITCH
+  priority: P2
+  hardware_state: 'ON'
+  ip_address: 192.168.1.3
+  software_state: GOOD
+  file_system_state: GOOD
+
+########################################
+# Links
+- itemType: LINK
+  id: '4'
+  name: link1
+  bandwidth: 1000
+  source: '1'
+  destination: '3'
+- itemType: LINK
+  id: '5'
+  name: link2
+  bandwidth: 1000
+  source: '3'
+  destination: '2'
diff --git a/tests/config/multidiscrete_obs_space_laydown_config.yaml b/tests/config/multidiscrete_obs_space_laydown_config.yaml
index dba1c6de..38438d6d 100644
--- a/tests/config/multidiscrete_obs_space_laydown_config.yaml
+++ b/tests/config/multidiscrete_obs_space_laydown_config.yaml
@@ -6,26 +6,63 @@
   steps: 5
 - itemType: PORTS
   portsList:
-  - port: '21'
+    - port: '80'
 - itemType: SERVICES
   serviceList:
-  - name: ftp
+    - name: TCP
+
+########################################
+# Nodes
 - itemType: NODE
   node_id: '1'
-  name: node
+  name: PC1
   node_class: SERVICE
   node_type: COMPUTER
-  priority: P1
+  priority: P5
   hardware_state: 'ON'
-  ip_address: 192.168.0.1
+  ip_address: 192.168.1.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
+    - name: TCP
+      port: '80'
+      state: GOOD
+- itemType: NODE
+  node_id: '2'
+  name: SERVER
+  node_class: SERVICE
+  node_type: SERVER
+  priority: P5
+  hardware_state: 'ON'
+  ip_address: 192.168.1.2
+  software_state: GOOD
+  file_system_state: GOOD
+  services:
+    - name: TCP
+      port: '80'
+      state: GOOD
+- itemType: NODE
+  node_id: '3'
+  name: SWITCH1
+  node_class: ACTIVE
+  node_type: SWITCH
+  priority: P2
+  hardware_state: 'ON'
+  ip_address: 192.168.1.3
+  software_state: GOOD
+  file_system_state: GOOD
+
+########################################
+# Links
+- itemType: LINK
+  id: '4'
+  name: link1
+  bandwidth: 1000
+  source: '1'
+  destination: '3'
+- itemType: LINK
+  id: '5'
+  name: link2
+  bandwidth: 1000
+  source: '3'
+  destination: '2'
diff --git a/tests/test_observation_space.py b/tests/test_observation_space.py
index 2f8e6a24..6a187761 100644
--- a/tests/test_observation_space.py
+++ b/tests/test_observation_space.py
@@ -12,6 +12,12 @@ def test_creating_env_with_box_obs():
     )
     env.update_environent_obs()
 
+    # we have three nodes and two links, with one service
+    # therefore the box observation space will have:
+    #   * 5 columns (four fixed and one for the service)
+    #   * 5 rows (3 nodes + 2 links)
+    assert env.env_obs.shape == (5, 5)
+
 
 def test_creating_env_with_multidiscrete_obs():
     """Try creating env with MultiDiscrete observation space."""
@@ -21,3 +27,8 @@ def test_creating_env_with_multidiscrete_obs():
         / "multidiscrete_obs_space_laydown_config.yaml",
     )
     env.update_environent_obs()
+
+    # we have three nodes and two links, with one service
+    # the nodes have hardware, OS, FS, and service, the links just have bandwidth,
+    # therefore we need 3*4 + 2 observations
+    assert env.env_obs.shape == (3 * 4 + 2,)