Apply suggestions from code review.

src/primaite/common/config_values_main.py

@@ -9,6 +9,7 @@ class ConfigValuesMain(object):
         """Init."""
         # Generic
         self.agent_identifier = ""  # the agent in use
+        self.observation_config = None  # observation space config
         self.num_episodes = 0  # number of episodes to train over
         self.num_steps = 0  # number of steps in an episode
         self.time_delay = 0  # delay between steps (ms) - applies to generic agents only

src/primaite/environment/observations.py

@@ -1,7 +1,7 @@
 """Module for handling configurable observation spaces in PrimAITE."""
 import logging
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, List, Tuple, Union
+from typing import TYPE_CHECKING, Dict, Final, List, Tuple, Union
 
 import numpy as np
 from gym import spaces
@@ -56,9 +56,9 @@ class NodeLinkTable(AbstractObservationComponent):
     ``(12, 7)``
     """
 
-    _FIXED_PARAMETERS = 4
+    _FIXED_PARAMETERS: int = 4
-    _MAX_VAL = 1_000_000
+    _MAX_VAL: int = 1_000_000
-    _DATA_TYPE = np.int64
+    _DATA_TYPE: type = np.int64
 
     def __init__(self, env: "Primaite"):
         super().__init__(env)
@@ -159,7 +159,7 @@ class NodeStatuses(AbstractObservationComponent):
     :type env: Primaite
     """
 
-    _DATA_TYPE = np.int64
+    _DATA_TYPE: type = np.int64
 
     def __init__(self, env: "Primaite"):
         super().__init__(env)
@@ -231,7 +231,7 @@ class LinkTrafficLevels(AbstractObservationComponent):
     :type quantisation_levels: int, optional
     """
 
-    _DATA_TYPE = np.int64
+    _DATA_TYPE: type = np.int64
 
     def __init__(
         self,
@@ -239,7 +239,14 @@ class LinkTrafficLevels(AbstractObservationComponent):
         combine_service_traffic: bool = False,
         quantisation_levels: int = 5,
     ):
-        assert quantisation_levels >= 3
+        if quantisation_levels < 3:
+            _msg = (
+                f"quantisation_levels must be 3 or more because the lowest and highest levels are "
+                f"reserved for 0% and 100% link utilisation, got {quantisation_levels} instead. "
+                f"Resetting to default value (5)"
+            )
+            _LOGGER.warning(_msg)
+            quantisation_levels = 5
 
         super().__init__(env)
 
@@ -296,7 +303,7 @@ class ObservationsHandler:
     Each component can also define further parameters to make them more flexible.
     """
 
-    registry = {
+    _REGISTRY: Final[Dict[str, type]] = {
         "NODE_LINK_TABLE": NodeLinkTable,
         "NODE_STATUSES": NodeStatuses,
         "LINK_TRAFFIC_LEVELS": LinkTrafficLevels,
@@ -384,7 +391,7 @@ class ObservationsHandler:
         for component_cfg in obs_space_config["components"]:
             # Figure out which class can instantiate the desired component
             comp_type = component_cfg["name"]
-            comp_class = cls.registry[comp_type]
+            comp_class = cls._REGISTRY[comp_type]
 
             # Create the component with options from the YAML
             options = component_cfg.get("options") or {}

src/primaite/environment/primaite_env.py

@@ -48,10 +48,10 @@ class Primaite(Env):
     """PRIMmary AI Training Evironment (Primaite) class."""
 
     #  Action Space contants
-    ACTION_SPACE_NODE_PROPERTY_VALUES = 5
+    ACTION_SPACE_NODE_PROPERTY_VALUES: int = 5
-    ACTION_SPACE_NODE_ACTION_VALUES = 4
+    ACTION_SPACE_NODE_ACTION_VALUES: int = 4
-    ACTION_SPACE_ACL_ACTION_VALUES = 3
+    ACTION_SPACE_ACL_ACTION_VALUES: int = 3
-    ACTION_SPACE_ACL_PERMISSION_VALUES = 2
+    ACTION_SPACE_ACL_PERMISSION_VALUES: int = 2
 
     def __init__(self, _config_values, _transaction_list):
         """
@@ -148,6 +148,8 @@ class Primaite(Env):
 
         # stores the observation config from the yaml, default is NODE_LINK_TABLE
         self.obs_config: dict = {"components": [{"name": "NODE_LINK_TABLE"}]}
+        if self.config_values.observation_config is not None:
+            self.obs_config = self.config_values.observation_config
 
         # Observation Handler manages the user-configurable observation space.
         # It will be initialised later.
@@ -690,9 +692,6 @@ class Primaite(Env):
             elif item["itemType"] == "ACTIONS":
                 # Get the action information
                 self.get_action_info(item)
-            elif item["itemType"] == "OBSERVATION_SPACE":
-                # Get the observation information
-                self.save_obs_config(item)
             elif item["itemType"] == "STEPS":
                 # Get the steps information
                 self.get_steps_info(item)

src/primaite/main.py

@@ -163,6 +163,10 @@ def load_config_values():
     try:
         # Generic
         config_values.agent_identifier = config_data["agentIdentifier"]
+        if "observationSpace" in config_data:
+            config_values.observation_config = config_data["observationSpace"]
+        else:
+            config_values.observation_config = None
         config_values.num_episodes = int(config_data["numEpisodes"])
         config_values.time_delay = int(config_data["timeDelay"])
         config_values.config_filename_use_case = (

tests/config/obs_tests/laydown.yaml

@@ -1,8 +1,5 @@
 - itemType: ACTIONS
   type: NODE
-- itemType: OBSERVATION_SPACE
-  components:
-    - name: NODE_STATUSES
 - itemType: STEPS
   steps: 5
 - itemType: PORTS
@@ -82,7 +79,7 @@
   id: '5'
   startStep: 0
   endStep: 5
-  load: 20
+  load: 999
   protocol: TCP
   port: '80'
   source: '1'

tests/config/obs_tests/laydown_with_LINK_TRAFFIC_LEVELS.yaml

@@ -1,110 +0,0 @@
-- itemType: ACTIONS
-  type: NODE
-- itemType: OBSERVATION_SPACE
-  components:
-    - name: LINK_TRAFFIC_LEVELS
-      options:
-        combine_service_traffic: false
-        quantisation_levels: 8
-- itemType: STEPS
-  steps: 5
-- itemType: PORTS
-  portsList:
-    - port: '80'
-    - port: '53'
-- itemType: SERVICES
-  serviceList:
-    - name: TCP
-    - name: UDP
-
-########################################
-# Nodes
-- itemType: NODE
-  node_id: '1'
-  name: PC1
-  node_class: SERVICE
-  node_type: COMPUTER
-  priority: P5
-  hardware_state: 'ON'
-  ip_address: 192.168.1.1
-  software_state: GOOD
-  file_system_state: GOOD
-  services:
-    - name: TCP
-      port: '80'
-      state: GOOD
-    - name: UDP
-      port: '53'
-      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
-    - name: UDP
-      port: '53'
-      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'
-
-#########################################
-# IERS
-- itemType: GREEN_IER
-  id: '5'
-  startStep: 0
-  endStep: 5
-  load: 999
-  protocol: TCP
-  port: '80'
-  source: '1'
-  destination: '2'
-  missionCriticality: 5
-
-#########################################
-# ACL Rules
-- itemType: ACL_RULE
-  id: '6'
-  permission: ALLOW
-  source: 192.168.1.1
-  destination: 192.168.1.2
-  protocol: TCP
-  port: 80
-- itemType: ACL_RULE
-  id: '7'
-  permission: ALLOW
-  source: 192.168.1.2
-  destination: 192.168.1.1
-  protocol: TCP
-  port: 80

tests/config/obs_tests/laydown_with_NODE_LINK_TABLE.yaml

@@ -1,77 +0,0 @@
-- itemType: ACTIONS
-  type: NODE
-- itemType: OBSERVATION_SPACE
-  components:
-    - name: NODE_LINK_TABLE
-- itemType: STEPS
-  steps: 5
-- itemType: PORTS
-  portsList:
-    - port: '80'
-    - port: '53'
-- itemType: SERVICES
-  serviceList:
-    - name: TCP
-    - name: UDP
-
-########################################
-# Nodes
-- itemType: NODE
-  node_id: '1'
-  name: PC1
-  node_class: SERVICE
-  node_type: COMPUTER
-  priority: P5
-  hardware_state: 'ON'
-  ip_address: 192.168.1.1
-  software_state: GOOD
-  file_system_state: GOOD
-  services:
-    - name: TCP
-      port: '80'
-      state: GOOD
-    - name: UDP
-      port: '53'
-      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
-    - name: UDP
-      port: '53'
-      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'

tests/config/obs_tests/laydown_without_obs_space.yaml

@@ -1,74 +0,0 @@
-- itemType: ACTIONS
-  type: NODE
-- itemType: STEPS
-  steps: 5
-- itemType: PORTS
-  portsList:
-    - port: '80'
-    - port: '53'
-- itemType: SERVICES
-  serviceList:
-    - name: TCP
-    - name: UDP
-
-########################################
-# Nodes
-- itemType: NODE
-  node_id: '1'
-  name: PC1
-  node_class: SERVICE
-  node_type: COMPUTER
-  priority: P5
-  hardware_state: 'ON'
-  ip_address: 192.168.1.1
-  software_state: GOOD
-  file_system_state: GOOD
-  services:
-    - name: TCP
-      port: '80'
-      state: GOOD
-    - name: UDP
-      port: '53'
-      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
-    - name: UDP
-      port: '53'
-      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'

tests/config/obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml

@@ -0,0 +1,96 @@
+# Main Config File
+
+# Generic config values
+# Choose one of these (dependent on Agent being trained)
+# "STABLE_BASELINES3_PPO"
+# "STABLE_BASELINES3_A2C"
+# "GENERIC"
+agentIdentifier: NONE
+# Number of episodes to run per session
+observationSpace:
+  components:
+    - name: LINK_TRAFFIC_LEVELS
+      options:
+        combine_service_traffic: false
+        quantisation_levels: 8
+
+numEpisodes: 1
+# Time delay between steps (for generic agents)
+timeDelay: 1
+# Filename of the scenario / laydown
+configFilename: one_node_states_on_off_lay_down_config.yaml
+# Type of session to be run (TRAINING or EVALUATION)
+sessionType: TRAINING
+# Determine whether to load an agent from file
+loadAgent: False
+# File path and file name of agent if you're loading one in
+agentLoadFile: C:\[Path]\[agent_saved_filename.zip]
+
+# Environment config values
+# The high value for the observation space
+observationSpaceHighValue: 1_000_000_000
+
+# Reward values
+# Generic
+allOk: 0
+# Node Hardware State
+offShouldBeOn: -10
+offShouldBeResetting: -5
+onShouldBeOff: -2
+onShouldBeResetting: -5
+resettingShouldBeOn: -5
+resettingShouldBeOff: -2
+resetting: -3
+# Node Software or Service State
+goodShouldBePatching: 2
+goodShouldBeCompromised: 5
+goodShouldBeOverwhelmed: 5
+patchingShouldBeGood: -5
+patchingShouldBeCompromised: 2
+patchingShouldBeOverwhelmed: 2
+patching: -3
+compromisedShouldBeGood: -20
+compromisedShouldBePatching: -20
+compromisedShouldBeOverwhelmed: -20
+compromised: -20
+overwhelmedShouldBeGood: -20
+overwhelmedShouldBePatching: -20
+overwhelmedShouldBeCompromised: -20
+overwhelmed: -20
+# Node File System State
+goodShouldBeRepairing: 2
+goodShouldBeRestoring: 2
+goodShouldBeCorrupt: 5
+goodShouldBeDestroyed: 10
+repairingShouldBeGood: -5
+repairingShouldBeRestoring: 2
+repairingShouldBeCorrupt: 2
+repairingShouldBeDestroyed: 0
+repairing: -3
+restoringShouldBeGood: -10
+restoringShouldBeRepairing: -2
+restoringShouldBeCorrupt: 1
+restoringShouldBeDestroyed: 2
+restoring: -6
+corruptShouldBeGood: -10
+corruptShouldBeRepairing: -10
+corruptShouldBeRestoring: -10
+corruptShouldBeDestroyed: 2
+corrupt: -10
+destroyedShouldBeGood: -20
+destroyedShouldBeRepairing: -20
+destroyedShouldBeRestoring: -20
+destroyedShouldBeCorrupt: -20
+destroyed: -20
+scanning: -2
+# IER status
+redIerRunning: -5
+greenIerBlocked: -10
+
+# Patching / Reset durations
+osPatchingDuration: 5            # The time taken to patch the OS
+nodeResetDuration: 5             # The time taken to reset a node (hardware)
+servicePatchingDuration: 5       # The time taken to patch a service
+fileSystemRepairingLimit: 5      # The time take to repair the file system
+fileSystemRestoringLimit: 5      # The time take to restore the file system
+fileSystemScanningLimit: 5       # The time taken to scan the file system

tests/config/obs_tests/main_config_NODE_LINK_TABLE.yaml

@@ -0,0 +1,93 @@
+# Main Config File
+
+# Generic config values
+# Choose one of these (dependent on Agent being trained)
+# "STABLE_BASELINES3_PPO"
+# "STABLE_BASELINES3_A2C"
+# "GENERIC"
+agentIdentifier: NONE
+# Number of episodes to run per session
+observationSpace:
+  components:
+    - name: NODE_LINK_TABLE
+
+numEpisodes: 1
+# Time delay between steps (for generic agents)
+timeDelay: 1
+# Filename of the scenario / laydown
+configFilename: one_node_states_on_off_lay_down_config.yaml
+# Type of session to be run (TRAINING or EVALUATION)
+sessionType: TRAINING
+# Determine whether to load an agent from file
+loadAgent: False
+# File path and file name of agent if you're loading one in
+agentLoadFile: C:\[Path]\[agent_saved_filename.zip]
+
+# Environment config values
+# The high value for the observation space
+observationSpaceHighValue: 1_000_000_000
+
+# Reward values
+# Generic
+allOk: 0
+# Node Hardware State
+offShouldBeOn: -10
+offShouldBeResetting: -5
+onShouldBeOff: -2
+onShouldBeResetting: -5
+resettingShouldBeOn: -5
+resettingShouldBeOff: -2
+resetting: -3
+# Node Software or Service State
+goodShouldBePatching: 2
+goodShouldBeCompromised: 5
+goodShouldBeOverwhelmed: 5
+patchingShouldBeGood: -5
+patchingShouldBeCompromised: 2
+patchingShouldBeOverwhelmed: 2
+patching: -3
+compromisedShouldBeGood: -20
+compromisedShouldBePatching: -20
+compromisedShouldBeOverwhelmed: -20
+compromised: -20
+overwhelmedShouldBeGood: -20
+overwhelmedShouldBePatching: -20
+overwhelmedShouldBeCompromised: -20
+overwhelmed: -20
+# Node File System State
+goodShouldBeRepairing: 2
+goodShouldBeRestoring: 2
+goodShouldBeCorrupt: 5
+goodShouldBeDestroyed: 10
+repairingShouldBeGood: -5
+repairingShouldBeRestoring: 2
+repairingShouldBeCorrupt: 2
+repairingShouldBeDestroyed: 0
+repairing: -3
+restoringShouldBeGood: -10
+restoringShouldBeRepairing: -2
+restoringShouldBeCorrupt: 1
+restoringShouldBeDestroyed: 2
+restoring: -6
+corruptShouldBeGood: -10
+corruptShouldBeRepairing: -10
+corruptShouldBeRestoring: -10
+corruptShouldBeDestroyed: 2
+corrupt: -10
+destroyedShouldBeGood: -20
+destroyedShouldBeRepairing: -20
+destroyedShouldBeRestoring: -20
+destroyedShouldBeCorrupt: -20
+destroyed: -20
+scanning: -2
+# IER status
+redIerRunning: -5
+greenIerBlocked: -10
+
+# Patching / Reset durations
+osPatchingDuration: 5            # The time taken to patch the OS
+nodeResetDuration: 5             # The time taken to reset a node (hardware)
+servicePatchingDuration: 5       # The time taken to patch a service
+fileSystemRepairingLimit: 5      # The time take to repair the file system
+fileSystemRestoringLimit: 5      # The time take to restore the file system
+fileSystemScanningLimit: 5       # The time taken to scan the file system

tests/config/obs_tests/main_config_NODE_STATUSES.yaml

@@ -0,0 +1,93 @@
+# Main Config File
+
+# Generic config values
+# Choose one of these (dependent on Agent being trained)
+# "STABLE_BASELINES3_PPO"
+# "STABLE_BASELINES3_A2C"
+# "GENERIC"
+agentIdentifier: NONE
+# Number of episodes to run per session
+observationSpace:
+  components:
+    - name: NODE_STATUSES
+
+numEpisodes: 1
+# Time delay between steps (for generic agents)
+timeDelay: 1
+# Filename of the scenario / laydown
+configFilename: one_node_states_on_off_lay_down_config.yaml
+# Type of session to be run (TRAINING or EVALUATION)
+sessionType: TRAINING
+# Determine whether to load an agent from file
+loadAgent: False
+# File path and file name of agent if you're loading one in
+agentLoadFile: C:\[Path]\[agent_saved_filename.zip]
+
+# Environment config values
+# The high value for the observation space
+observationSpaceHighValue: 1_000_000_000
+
+# Reward values
+# Generic
+allOk: 0
+# Node Hardware State
+offShouldBeOn: -10
+offShouldBeResetting: -5
+onShouldBeOff: -2
+onShouldBeResetting: -5
+resettingShouldBeOn: -5
+resettingShouldBeOff: -2
+resetting: -3
+# Node Software or Service State
+goodShouldBePatching: 2
+goodShouldBeCompromised: 5
+goodShouldBeOverwhelmed: 5
+patchingShouldBeGood: -5
+patchingShouldBeCompromised: 2
+patchingShouldBeOverwhelmed: 2
+patching: -3
+compromisedShouldBeGood: -20
+compromisedShouldBePatching: -20
+compromisedShouldBeOverwhelmed: -20
+compromised: -20
+overwhelmedShouldBeGood: -20
+overwhelmedShouldBePatching: -20
+overwhelmedShouldBeCompromised: -20
+overwhelmed: -20
+# Node File System State
+goodShouldBeRepairing: 2
+goodShouldBeRestoring: 2
+goodShouldBeCorrupt: 5
+goodShouldBeDestroyed: 10
+repairingShouldBeGood: -5
+repairingShouldBeRestoring: 2
+repairingShouldBeCorrupt: 2
+repairingShouldBeDestroyed: 0
+repairing: -3
+restoringShouldBeGood: -10
+restoringShouldBeRepairing: -2
+restoringShouldBeCorrupt: 1
+restoringShouldBeDestroyed: 2
+restoring: -6
+corruptShouldBeGood: -10
+corruptShouldBeRepairing: -10
+corruptShouldBeRestoring: -10
+corruptShouldBeDestroyed: 2
+corrupt: -10
+destroyedShouldBeGood: -20
+destroyedShouldBeRepairing: -20
+destroyedShouldBeRestoring: -20
+destroyedShouldBeCorrupt: -20
+destroyed: -20
+scanning: -2
+# IER status
+redIerRunning: -5
+greenIerBlocked: -10
+
+# Patching / Reset durations
+osPatchingDuration: 5            # The time taken to patch the OS
+nodeResetDuration: 5             # The time taken to reset a node (hardware)
+servicePatchingDuration: 5       # The time taken to patch a service
+fileSystemRepairingLimit: 5      # The time take to repair the file system
+fileSystemRestoringLimit: 5      # The time take to restore the file system
+fileSystemScanningLimit: 5       # The time taken to scan the file system

tests/config/obs_tests/main_config_without_obs.yaml

@@ -21,7 +21,7 @@ agentLoadFile: C:\[Path]\[agent_saved_filename.zip]
 
 # Environment config values
 # The high value for the observation space
-observationSpaceHighValue: 1000000000
+observationSpaceHighValue: 1_000_000_000
 
 # Reward values
 # Generic

tests/conftest.py

@@ -19,6 +19,10 @@ def _get_primaite_env_from_config(
 
     def load_config_values():
         config_values.agent_identifier = config_data["agentIdentifier"]
+        if "observationSpace" in config_data:
+            config_values.observation_config = config_data["observationSpace"]
+        else:
+            config_values.observation_config = None
         config_values.num_episodes = int(config_data["numEpisodes"])
         config_values.time_delay = int(config_data["timeDelay"])
         config_values.config_filename_use_case = lay_down_config_path

tests/test_observation_space.py

@@ -27,9 +27,8 @@ def env(request):
 
 @pytest.mark.env_config_paths(
     dict(
-        main_config_path=TEST_CONFIG_ROOT / "one_node_states_on_off_main_config.yaml",
+        main_config_path=TEST_CONFIG_ROOT / "obs_tests/main_config_without_obs.yaml",
-        lay_down_config_path=TEST_CONFIG_ROOT
+        lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
-        / "obs_tests/laydown_without_obs_space.yaml",
     )
 )
 def test_default_obs_space(env: Primaite):
@@ -44,9 +43,8 @@ def test_default_obs_space(env: Primaite):
 
 @pytest.mark.env_config_paths(
     dict(
-        main_config_path=TEST_CONFIG_ROOT / "one_node_states_on_off_main_config.yaml",
+        main_config_path=TEST_CONFIG_ROOT / "obs_tests/main_config_without_obs.yaml",
-        lay_down_config_path=TEST_CONFIG_ROOT
+        lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
-        / "obs_tests/laydown_without_obs_space.yaml",
     )
 )
 def test_registering_components(env: Primaite):
@@ -61,9 +59,9 @@ def test_registering_components(env: Primaite):
 
 @pytest.mark.env_config_paths(
     dict(
-        main_config_path=TEST_CONFIG_ROOT / "obs_tests/main_config_no_agent.yaml",
+        main_config_path=TEST_CONFIG_ROOT
-        lay_down_config_path=TEST_CONFIG_ROOT
+        / "obs_tests/main_config_NODE_LINK_TABLE.yaml",
-        / "obs_tests/laydown_with_NODE_LINK_TABLE.yaml",
+        lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
     )
 )
 class TestNodeLinkTable:
@@ -92,17 +90,17 @@ class TestNodeLinkTable:
             * Node 1:
                 * 1 (id)
                 * 1 (good hardware state)
-                * 1 (good OS state)
+                * 3 (compromised OS state)
                 * 1 (good file system state)
-                * 1 (good service1 state)
+                * 1 (good TCP state)
-                * 1 (good service2 state)
+                * 1 (good UDP state)
             * Node 2:
                 * 2 (id)
                 * 1 (good hardware state)
                 * 1 (good OS state)
                 * 1 (good file system state)
-                * 1 (good service1 state)
+                * 1 (good TCP state)
-                * 1 (good service2 state)
+                * 4 (overwhelmed UDP state)
             * Node 3 (active node):
                 * 3 (id)
                 * 1 (good hardware state)
@@ -115,14 +113,14 @@ class TestNodeLinkTable:
                 * 0 (n/a hardware state)
                 * 0 (n/a OS state)
                 * 0 (n/a file system state)
-                * 0 (no traffic for service1)
+                * 999 (999 traffic for service1)
                 * 0 (no traffic for service2)
             * Link 2:
                 * 5 (id)
                 * 0 (good hardware state)
                 * 0 (good OS state)
                 * 0 (good file system state)
-                * 0 (no traffic service1)
+                * 999 (999 traffic service1)
                 * 0 (no traffic for service2)
         """
         act = np.asarray([0, 0, 0, 0])
@@ -131,20 +129,19 @@ class TestNodeLinkTable:
         assert np.array_equal(
             obs,
             [
-                [1, 1, 1, 1, 1, 1],
+                [1, 1, 3, 1, 1, 1],
-                [2, 1, 1, 1, 1, 1],
+                [2, 1, 1, 1, 1, 4],
                 [3, 1, 1, 1, 0, 0],
-                [4, 0, 0, 0, 0, 0],
+                [4, 0, 0, 0, 999, 0],
-                [5, 0, 0, 0, 0, 0],
+                [5, 0, 0, 0, 999, 0],
             ],
         )
 
 
 @pytest.mark.env_config_paths(
     dict(
-        main_config_path=TEST_CONFIG_ROOT / "obs_tests/main_config_no_agent.yaml",
+        main_config_path=TEST_CONFIG_ROOT / "obs_tests/main_config_NODE_STATUSES.yaml",
-        lay_down_config_path=TEST_CONFIG_ROOT
+        lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
-        / "obs_tests/laydown_with_NODE_STATUSES.yaml",
     )
 )
 class TestNodeStatuses:
@@ -188,9 +185,9 @@ class TestNodeStatuses:
 
 @pytest.mark.env_config_paths(
     dict(
-        main_config_path=TEST_CONFIG_ROOT / "obs_tests/main_config_no_agent.yaml",
+        main_config_path=TEST_CONFIG_ROOT
-        lay_down_config_path=TEST_CONFIG_ROOT
+        / "obs_tests/main_config_LINK_TRAFFIC_LEVELS.yaml",
-        / "obs_tests/laydown_with_LINK_TRAFFIC_LEVELS.yaml",
+        lay_down_config_path=TEST_CONFIG_ROOT / "obs_tests/laydown.yaml",
     )
 )
 class TestLinkTrafficLevels: