diff --git a/docs/source/simulation_components/system/software.rst b/docs/source/simulation_components/system/software.rst
index 1e5a0b6b..cd6b0aa3 100644
--- a/docs/source/simulation_components/system/software.rst
+++ b/docs/source/simulation_components/system/software.rst
@@ -28,7 +28,7 @@ See :ref:`Node Start up and Shut down`
 
     node.software_manager.install(WebServer)
 
-    web_server: WebServer = node.software_manager.software["WebServer"]
+    web_server: WebServer = node.software_manager.software.get("WebServer")
     assert web_server.operating_state is ServiceOperatingState.RUNNING # service is immediately ran after install
 
     node.power_off()
diff --git a/src/primaite/simulator/domain/account.py b/src/primaite/simulator/domain/account.py
index 1402a474..d9dad06a 100644
--- a/src/primaite/simulator/domain/account.py
+++ b/src/primaite/simulator/domain/account.py
@@ -72,7 +72,7 @@ class Account(SimComponent):
                 "num_group_changes": self.num_group_changes,
                 "username": self.username,
                 "password": self.password,
-                "account_type": self.account_type.name,
+                "account_type": self.account_type.value,
                 "enabled": self.enabled,
             }
         )
diff --git a/src/primaite/simulator/network/networks.py b/src/primaite/simulator/network/networks.py
index 0b6fe8d4..4cd9c8d3 100644
--- a/src/primaite/simulator/network/networks.py
+++ b/src/primaite/simulator/network/networks.py
@@ -147,7 +147,7 @@ def arcd_uc2_network() -> Network:
     client_1.power_on()
     network.connect(endpoint_b=client_1.ethernet_port[1], endpoint_a=switch_2.switch_ports[1])
     client_1.software_manager.install(DataManipulationBot)
-    db_manipulation_bot: DataManipulationBot = client_1.software_manager.software["DataManipulationBot"]
+    db_manipulation_bot: DataManipulationBot = client_1.software_manager.software.get("DataManipulationBot")
     db_manipulation_bot.configure(
         server_ip_address=IPv4Address("192.168.1.14"),
         payload="DELETE",
@@ -165,7 +165,7 @@ def arcd_uc2_network() -> Network:
         operating_state=NodeOperatingState.ON,
     )
     client_2.power_on()
-    web_browser = client_2.software_manager.software["WebBrowser"]
+    web_browser = client_2.software_manager.software.get("WebBrowser")
     web_browser.target_url = "http://arcd.com/users/"
     network.connect(endpoint_b=client_2.ethernet_port[1], endpoint_a=switch_2.switch_ports[2])
 
@@ -249,7 +249,7 @@ def arcd_uc2_network() -> Network:
         # noqa
     ]
     database_server.software_manager.install(DatabaseService)
-    database_service: DatabaseService = database_server.software_manager.software["DatabaseService"]  # noqa
+    database_service: DatabaseService = database_server.software_manager.software.get("DatabaseService")  # noqa
     database_service.start()
     database_service.configure_backup(backup_server=IPv4Address("192.168.1.16"))
     database_service._process_sql(ddl, None)  # noqa
@@ -268,7 +268,7 @@ def arcd_uc2_network() -> Network:
     web_server.power_on()
     web_server.software_manager.install(DatabaseClient)
 
-    database_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
+    database_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
     database_client.configure(server_ip_address=IPv4Address("192.168.1.14"))
     network.connect(endpoint_b=web_server.ethernet_port[1], endpoint_a=switch_1.switch_ports[2])
     database_client.run()
@@ -277,7 +277,7 @@ def arcd_uc2_network() -> Network:
     web_server.software_manager.install(WebServer)
 
     # register the web_server to a domain
-    dns_server_service: DNSServer = domain_controller.software_manager.software["DNSServer"]  # noqa
+    dns_server_service: DNSServer = domain_controller.software_manager.software.get("DNSServer")  # noqa
     dns_server_service.dns_register("arcd.com", web_server.ip_address)
 
     # Backup Server
diff --git a/src/primaite/simulator/system/applications/database_client.py b/src/primaite/simulator/system/applications/database_client.py
index 7b63d26e..8c43c0b7 100644
--- a/src/primaite/simulator/system/applications/database_client.py
+++ b/src/primaite/simulator/system/applications/database_client.py
@@ -107,7 +107,7 @@ class DatabaseClient(Application):
 
     def disconnect(self):
         """Disconnect from the Database Service."""
-        if self.connected and self.operating_state.RUNNING:
+        if self.connected and self.operating_state is ApplicationOperatingState.RUNNING:
             software_manager: SoftwareManager = self.software_manager
             software_manager.send_payload_to_session_manager(
                 payload={"type": "disconnect"}, dest_ip_address=self.server_ip_address, dest_port=self.port
@@ -186,6 +186,9 @@ class DatabaseClient(Application):
         :param session_id: The session id the payload relates to.
         :return: True.
         """
+        if not self._can_perform_action():
+            return False
+
         if isinstance(payload, dict) and payload.get("type"):
             if payload["type"] == "connect_response":
                 self.connected = payload["response"] == True
diff --git a/src/primaite/simulator/system/applications/web_browser.py b/src/primaite/simulator/system/applications/web_browser.py
index 1531314d..7533f6f3 100644
--- a/src/primaite/simulator/system/applications/web_browser.py
+++ b/src/primaite/simulator/system/applications/web_browser.py
@@ -99,7 +99,7 @@ class WebBrowser(Application):
             return False
 
         # get the IP address of the domain name via DNS
-        dns_client: DNSClient = self.software_manager.software["DNSClient"]
+        dns_client: DNSClient = self.software_manager.software.get("DNSClient")
         domain_exists = dns_client.check_domain_exists(target_domain=parsed_url.hostname)
 
         # if domain does not exist, the request fails
diff --git a/src/primaite/simulator/system/services/database/database_service.py b/src/primaite/simulator/system/services/database/database_service.py
index f9621ba5..6a7c80ca 100644
--- a/src/primaite/simulator/system/services/database/database_service.py
+++ b/src/primaite/simulator/system/services/database/database_service.py
@@ -80,7 +80,7 @@ class DatabaseService(Service):
             return False
 
         software_manager: SoftwareManager = self.software_manager
-        ftp_client_service: FTPClient = software_manager.software["FTPClient"]
+        ftp_client_service: FTPClient = software_manager.software.get("FTPClient")
 
         # send backup copy of database file to FTP server
         response = ftp_client_service.send_file(
@@ -104,7 +104,7 @@ class DatabaseService(Service):
             return False
 
         software_manager: SoftwareManager = self.software_manager
-        ftp_client_service: FTPClient = software_manager.software["FTPClient"]
+        ftp_client_service: FTPClient = software_manager.software.get("FTPClient")
 
         # retrieve backup file from backup server
         response = ftp_client_service.request_file(
diff --git a/src/primaite/simulator/system/services/service.py b/src/primaite/simulator/system/services/service.py
index 6d6cda86..e60b7700 100644
--- a/src/primaite/simulator/system/services/service.py
+++ b/src/primaite/simulator/system/services/service.py
@@ -109,8 +109,8 @@ class Service(IOSoftware):
         """
         state = super().describe_state()
         state["operating_state"] = self.operating_state.value
-        state["health_state_actual"] = self.health_state_actual
-        state["health_state_visible"] = self.health_state_visible
+        state["health_state_actual"] = self.health_state_actual.value
+        state["health_state_visible"] = self.health_state_visible.value
         return state
 
     def stop(self) -> None:
diff --git a/src/primaite/simulator/system/services/web_server/web_server.py b/src/primaite/simulator/system/services/web_server/web_server.py
index bff29a47..e63b875a 100644
--- a/src/primaite/simulator/system/services/web_server/web_server.py
+++ b/src/primaite/simulator/system/services/web_server/web_server.py
@@ -120,7 +120,7 @@ class WebServer(Service):
 
             if path.startswith("users"):
                 # get data from DatabaseServer
-                db_client: DatabaseClient = self.software_manager.software["DatabaseClient"]
+                db_client: DatabaseClient = self.software_manager.software.get("DatabaseClient")
                 # get all users
                 if db_client.query("SELECT"):
                     # query succeeded
diff --git a/tests/conftest.py b/tests/conftest.py
index 55db53c5..c81e4b98 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -228,7 +228,7 @@ def example_network() -> Network:
         default_gateway="192.168.1.1",
         operating_state=NodeOperatingState.ON,
     )
-    network.connect(endpoint_b=server_2.ethernet_port[1], endpoint_a=switch_1.switch_ports[3])
+    network.connect(endpoint_b=server_2.ethernet_port[1], endpoint_a=switch_1.switch_ports[2])
 
     router_1.acl.add_rule(action=ACLAction.PERMIT, src_port=Port.ARP, dst_port=Port.ARP, position=22)
     router_1.acl.add_rule(action=ACLAction.PERMIT, protocol=IPProtocol.ICMP, position=23)
diff --git a/tests/e2e_integration_tests/environments/test_sb3_environment.py b/tests/e2e_integration_tests/environments/test_sb3_environment.py
index c1c028a2..91cf5c1e 100644
--- a/tests/e2e_integration_tests/environments/test_sb3_environment.py
+++ b/tests/e2e_integration_tests/environments/test_sb3_environment.py
@@ -11,7 +11,7 @@ from primaite.game.game import PrimaiteGame
 from primaite.session.environment import PrimaiteGymEnv
 
 
-@pytest.mark.skip(reason="no way of currently testing this")
+# @pytest.mark.skip(reason="no way of currently testing this")
 def test_sb3_compatibility():
     """Test that the Gymnasium environment can be used with an SB3 agent."""
     with open(example_config_path(), "r") as f:
diff --git a/tests/e2e_integration_tests/test_primaite_session.py b/tests/e2e_integration_tests/test_primaite_session.py
index ed10ca24..7785e4ae 100644
--- a/tests/e2e_integration_tests/test_primaite_session.py
+++ b/tests/e2e_integration_tests/test_primaite_session.py
@@ -11,7 +11,7 @@ MISCONFIGURED_PATH = TEST_ASSETS_ROOT / "configs/bad_primaite_session.yaml"
 MULTI_AGENT_PATH = TEST_ASSETS_ROOT / "configs/multi_agent_session.yaml"
 
 
-@pytest.mark.skip(reason="no way of currently testing this")
+# @pytest.mark.skip(reason="no way of currently testing this")
 class TestPrimaiteSession:
     @pytest.mark.parametrize("temp_primaite_session", [[CFG_PATH]], indirect=True)
     def test_creating_session(self, temp_primaite_session):
diff --git a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py
index 81bbfc96..0dc2c031 100644
--- a/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py
+++ b/tests/e2e_integration_tests/test_uc2_data_manipulation_scenario.py
@@ -8,13 +8,13 @@ from primaite.simulator.system.services.red_services.data_manipulation_bot impor
 def test_data_manipulation(uc2_network):
     """Tests the UC2 data manipulation scenario end-to-end. Is a work in progress."""
     client_1: Computer = uc2_network.get_node_by_hostname("client_1")
-    db_manipulation_bot: DataManipulationBot = client_1.software_manager.software["DataManipulationBot"]
+    db_manipulation_bot: DataManipulationBot = client_1.software_manager.software.get("DataManipulationBot")
 
     database_server: Server = uc2_network.get_node_by_hostname("database_server")
-    db_service: DatabaseService = database_server.software_manager.software["DatabaseService"]
+    db_service: DatabaseService = database_server.software_manager.software.get("DatabaseService")
 
     web_server: Server = uc2_network.get_node_by_hostname("web_server")
-    db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
+    db_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
 
     db_service.backup_database()
 
diff --git a/tests/integration_tests/system/test_application_on_node.py b/tests/integration_tests/system/test_application_on_node.py
index cce586da..46be5e55 100644
--- a/tests/integration_tests/system/test_application_on_node.py
+++ b/tests/integration_tests/system/test_application_on_node.py
@@ -18,7 +18,7 @@ def populated_node(application_class) -> Tuple[Application, Computer]:
     )
     computer.software_manager.install(application_class)
 
-    app = computer.software_manager.software["TestApplication"]
+    app = computer.software_manager.software.get("TestApplication")
     app.run()
 
     return app, computer
@@ -35,7 +35,7 @@ def test_service_on_offline_node(application_class):
     )
     computer.software_manager.install(application_class)
 
-    app: Application = computer.software_manager.software["TestApplication"]
+    app: Application = computer.software_manager.software.get("TestApplication")
 
     computer.power_off()
 
diff --git a/tests/integration_tests/system/test_database_on_node.py b/tests/integration_tests/system/test_database_on_node.py
index ef2b2956..98c8c87b 100644
--- a/tests/integration_tests/system/test_database_on_node.py
+++ b/tests/integration_tests/system/test_database_on_node.py
@@ -10,10 +10,10 @@ from primaite.simulator.system.services.service import ServiceOperatingState
 
 def test_database_client_server_connection(uc2_network):
     web_server: Server = uc2_network.get_node_by_hostname("web_server")
-    db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
+    db_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
 
     db_server: Server = uc2_network.get_node_by_hostname("database_server")
-    db_service: DatabaseService = db_server.software_manager.software["DatabaseService"]
+    db_service: DatabaseService = db_server.software_manager.software.get("DatabaseService")
 
     assert len(db_service.connections) == 1
 
@@ -23,10 +23,10 @@ def test_database_client_server_connection(uc2_network):
 
 def test_database_client_server_correct_password(uc2_network):
     web_server: Server = uc2_network.get_node_by_hostname("web_server")
-    db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
+    db_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
 
     db_server: Server = uc2_network.get_node_by_hostname("database_server")
-    db_service: DatabaseService = db_server.software_manager.software["DatabaseService"]
+    db_service: DatabaseService = db_server.software_manager.software.get("DatabaseService")
 
     db_client.disconnect()
 
@@ -40,10 +40,10 @@ def test_database_client_server_correct_password(uc2_network):
 
 def test_database_client_server_incorrect_password(uc2_network):
     web_server: Server = uc2_network.get_node_by_hostname("web_server")
-    db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
+    db_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
 
     db_server: Server = uc2_network.get_node_by_hostname("database_server")
-    db_service: DatabaseService = db_server.software_manager.software["DatabaseService"]
+    db_service: DatabaseService = db_server.software_manager.software.get("DatabaseService")
 
     db_client.disconnect()
     db_client.configure(server_ip_address=IPv4Address("192.168.1.14"), server_password="54321")
@@ -56,7 +56,7 @@ def test_database_client_server_incorrect_password(uc2_network):
 def test_database_client_query(uc2_network):
     """Tests DB query across the network returns HTTP status 200 and date."""
     web_server: Server = uc2_network.get_node_by_hostname("web_server")
-    db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
+    db_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
 
     assert db_client.connected
 
@@ -66,13 +66,13 @@ def test_database_client_query(uc2_network):
 def test_create_database_backup(uc2_network):
     """Run the backup_database method and check if the FTP server has the relevant file."""
     db_server: Server = uc2_network.get_node_by_hostname("database_server")
-    db_service: DatabaseService = db_server.software_manager.software["DatabaseService"]
+    db_service: DatabaseService = db_server.software_manager.software.get("DatabaseService")
 
     # back up should be created
     assert db_service.backup_database() is True
 
     backup_server: Server = uc2_network.get_node_by_hostname("backup_server")
-    ftp_server: FTPServer = backup_server.software_manager.software["FTPServer"]
+    ftp_server: FTPServer = backup_server.software_manager.software.get("FTPServer")
 
     # backup file should exist in the backup server
     assert ftp_server.file_system.get_file(folder_name=db_service.uuid, file_name="database.db") is not None
@@ -81,7 +81,7 @@ def test_create_database_backup(uc2_network):
 def test_restore_backup(uc2_network):
     """Run the restore_backup method and check if the backup is properly restored."""
     db_server: Server = uc2_network.get_node_by_hostname("database_server")
-    db_service: DatabaseService = db_server.software_manager.software["DatabaseService"]
+    db_service: DatabaseService = db_server.software_manager.software.get("DatabaseService")
 
     # create a back up
     assert db_service.backup_database() is True
@@ -100,13 +100,13 @@ def test_restore_backup(uc2_network):
 def test_database_client_cannot_query_offline_database_server(uc2_network):
     """Tests DB query across the network returns HTTP status 404 when db server is offline."""
     db_server: Server = uc2_network.get_node_by_hostname("database_server")
-    db_service: DatabaseService = db_server.software_manager.software["DatabaseService"]
+    db_service: DatabaseService = db_server.software_manager.software.get("DatabaseService")
 
     assert db_server.operating_state is NodeOperatingState.ON
     assert db_service.operating_state is ServiceOperatingState.RUNNING
 
     web_server: Server = uc2_network.get_node_by_hostname("web_server")
-    db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
+    db_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
     assert db_client.connected
 
     assert db_client.query("SELECT") is True
diff --git a/tests/integration_tests/system/test_dns_client_server.py b/tests/integration_tests/system/test_dns_client_server.py
index 70657112..a54bf23f 100644
--- a/tests/integration_tests/system/test_dns_client_server.py
+++ b/tests/integration_tests/system/test_dns_client_server.py
@@ -17,14 +17,14 @@ def dns_client_and_dns_server(client_server) -> Tuple[DNSClient, Computer, DNSSe
 
     # Install DNS Client on computer
     computer.software_manager.install(DNSClient)
-    dns_client: DNSClient = computer.software_manager.software["DNSClient"]
+    dns_client: DNSClient = computer.software_manager.software.get("DNSClient")
     dns_client.start()
     # set server as DNS Server
     dns_client.dns_server = IPv4Address(server.nics.get(next(iter(server.nics))).ip_address)
 
     # Install DNS Server on server
     server.software_manager.install(DNSServer)
-    dns_server: DNSServer = server.software_manager.software["DNSServer"]
+    dns_server: DNSServer = server.software_manager.software.get("DNSServer")
     dns_server.start()
     # register arcd.com as a domain
     dns_server.dns_register(
diff --git a/tests/integration_tests/system/test_ftp_client_server.py b/tests/integration_tests/system/test_ftp_client_server.py
index 32ea7f2b..1a6a8f41 100644
--- a/tests/integration_tests/system/test_ftp_client_server.py
+++ b/tests/integration_tests/system/test_ftp_client_server.py
@@ -16,12 +16,12 @@ def ftp_client_and_ftp_server(client_server) -> Tuple[FTPClient, Computer, FTPSe
 
     # Install FTP Client service on computer
     computer.software_manager.install(FTPClient)
-    ftp_client: FTPClient = computer.software_manager.software["FTPClient"]
+    ftp_client: FTPClient = computer.software_manager.software.get("FTPClient")
     ftp_client.start()
 
     # Install FTP Server service on server
     server.software_manager.install(FTPServer)
-    ftp_server: FTPServer = server.software_manager.software["FTPServer"]
+    ftp_server: FTPServer = server.software_manager.software.get("FTPServer")
     ftp_server.start()
 
     return ftp_client, computer, ftp_server, server
diff --git a/tests/integration_tests/system/test_service_on_node.py b/tests/integration_tests/system/test_service_on_node.py
index 9480c358..aab1e4da 100644
--- a/tests/integration_tests/system/test_service_on_node.py
+++ b/tests/integration_tests/system/test_service_on_node.py
@@ -17,7 +17,7 @@ def populated_node(
     )
     server.software_manager.install(service_class)
 
-    service = server.software_manager.software["TestService"]
+    service = server.software_manager.software.get("TestService")
     service.start()
 
     return server, service
@@ -34,7 +34,7 @@ def test_service_on_offline_node(service_class):
     )
     computer.software_manager.install(service_class)
 
-    service: Service = computer.software_manager.software["TestService"]
+    service: Service = computer.software_manager.software.get("TestService")
 
     computer.power_off()
 
diff --git a/tests/integration_tests/system/test_web_client_server.py b/tests/integration_tests/system/test_web_client_server.py
index 41982805..b3d2e891 100644
--- a/tests/integration_tests/system/test_web_client_server.py
+++ b/tests/integration_tests/system/test_web_client_server.py
@@ -19,23 +19,23 @@ def web_client_and_web_server(client_server) -> Tuple[WebBrowser, Computer, WebS
 
     # Install Web Browser on computer
     computer.software_manager.install(WebBrowser)
-    web_browser: WebBrowser = computer.software_manager.software["WebBrowser"]
+    web_browser: WebBrowser = computer.software_manager.software.get("WebBrowser")
     web_browser.run()
 
     # Install DNS Client service on computer
     computer.software_manager.install(DNSClient)
-    dns_client: DNSClient = computer.software_manager.software["DNSClient"]
+    dns_client: DNSClient = computer.software_manager.software.get("DNSClient")
     # set dns server
     dns_client.dns_server = server.nics[next(iter(server.nics))].ip_address
 
     # Install Web Server service on server
     server.software_manager.install(WebServer)
-    web_server_service: WebServer = server.software_manager.software["WebServer"]
+    web_server_service: WebServer = server.software_manager.software.get("WebServer")
     web_server_service.start()
 
     # Install DNS Server service on server
     server.software_manager.install(DNSServer)
-    dns_server: DNSServer = server.software_manager.software["DNSServer"]
+    dns_server: DNSServer = server.software_manager.software.get("DNSServer")
     # register arcd.com to DNS
     dns_server.dns_register(domain_name="arcd.com", domain_ip_address=server.nics[next(iter(server.nics))].ip_address)
 
diff --git a/tests/integration_tests/system/test_web_client_server_and_database.py b/tests/integration_tests/system/test_web_client_server_and_database.py
index d7b5603d..17458968 100644
--- a/tests/integration_tests/system/test_web_client_server_and_database.py
+++ b/tests/integration_tests/system/test_web_client_server_and_database.py
@@ -60,28 +60,28 @@ def web_client_web_server_database() -> Tuple[Computer, Server, Server]:
 
     # Install DatabaseService on db server
     db_server.software_manager.install(DatabaseService)
-    db_service: DatabaseService = db_server.software_manager.software["DatabaseService"]
+    db_service: DatabaseService = db_server.software_manager.software.get("DatabaseService")
     db_service.start()
 
     # Install Web Browser on computer
     computer.software_manager.install(WebBrowser)
-    web_browser: WebBrowser = computer.software_manager.software["WebBrowser"]
+    web_browser: WebBrowser = computer.software_manager.software.get("WebBrowser")
     web_browser.run()
 
     # Install DNS Client service on computer
     computer.software_manager.install(DNSClient)
-    dns_client: DNSClient = computer.software_manager.software["DNSClient"]
+    dns_client: DNSClient = computer.software_manager.software.get("DNSClient")
     # set dns server
     dns_client.dns_server = web_server.nics[next(iter(web_server.nics))].ip_address
 
     # Install Web Server service on web server
     web_server.software_manager.install(WebServer)
-    web_server_service: WebServer = web_server.software_manager.software["WebServer"]
+    web_server_service: WebServer = web_server.software_manager.software.get("WebServer")
     web_server_service.start()
 
     # Install DNS Server service on web server
     web_server.software_manager.install(DNSServer)
-    dns_server: DNSServer = web_server.software_manager.software["DNSServer"]
+    dns_server: DNSServer = web_server.software_manager.software.get("DNSServer")
     # register arcd.com to DNS
     dns_server.dns_register(
         domain_name="arcd.com", domain_ip_address=web_server.nics[next(iter(web_server.nics))].ip_address
@@ -89,7 +89,7 @@ def web_client_web_server_database() -> Tuple[Computer, Server, Server]:
 
     # Install DatabaseClient service on web server
     web_server.software_manager.install(DatabaseClient)
-    db_client: DatabaseClient = web_server.software_manager.software["DatabaseClient"]
+    db_client: DatabaseClient = web_server.software_manager.software.get("DatabaseClient")
     db_client.server_ip_address = IPv4Address(db_server_nic.ip_address)  # set IP address of Database Server
     db_client.run()
     assert db_client.connect()
@@ -101,6 +101,6 @@ def web_client_web_server_database() -> Tuple[Computer, Server, Server]:
 def test_web_client_requests_users(web_client_web_server_database):
     computer, web_server, db_server = web_client_web_server_database
 
-    web_browser: WebBrowser = computer.software_manager.software["WebBrowser"]
+    web_browser: WebBrowser = computer.software_manager.software.get("WebBrowser")
 
     web_browser.get_webpage()
diff --git a/tests/unit_tests/_primaite/_simulator/_domain/test_account.py b/tests/unit_tests/_primaite/_simulator/_domain/test_account.py
index 96c34996..01ad3871 100644
--- a/tests/unit_tests/_primaite/_simulator/_domain/test_account.py
+++ b/tests/unit_tests/_primaite/_simulator/_domain/test_account.py
@@ -1,18 +1,140 @@
 """Test the account module of the simulator."""
+import pytest
+
 from primaite.simulator.domain.account import Account, AccountType
 
 
-def test_account_serialise():
+@pytest.fixture(scope="function")
+def account() -> Account:
+    acct = Account(username="Jake", password="totally_hashed_password", account_type=AccountType.USER)
+    acct.set_original_state()
+    return acct
+
+
+def test_original_state(account):
+    """Test the original state - see if it resets properly"""
+    account.log_on()
+    account.log_off()
+    account.disable()
+
+    state = account.describe_state()
+    assert state["num_logons"] is 1
+    assert state["num_logoffs"] is 1
+    assert state["num_group_changes"] is 0
+    assert state["username"] is "Jake"
+    assert state["password"] is "totally_hashed_password"
+    assert state["account_type"] is AccountType.USER.value
+    assert state["enabled"] is False
+
+    account.reset_component_for_episode(episode=1)
+    state = account.describe_state()
+    assert state["num_logons"] is 0
+    assert state["num_logoffs"] is 0
+    assert state["num_group_changes"] is 0
+    assert state["username"] is "Jake"
+    assert state["password"] is "totally_hashed_password"
+    assert state["account_type"] is AccountType.USER.value
+    assert state["enabled"] is True
+
+    account.log_on()
+    account.log_off()
+    account.disable()
+    account.set_original_state()
+
+    account.log_on()
+    state = account.describe_state()
+    assert state["num_logons"] is 2
+
+    account.reset_component_for_episode(episode=2)
+    state = account.describe_state()
+    assert state["num_logons"] is 1
+    assert state["num_logoffs"] is 1
+    assert state["num_group_changes"] is 0
+    assert state["username"] is "Jake"
+    assert state["password"] is "totally_hashed_password"
+    assert state["account_type"] is AccountType.USER.value
+    assert state["enabled"] is False
+
+
+def test_enable(account):
+    """Should enable the account."""
+    account.enabled = False
+    account.enable()
+    assert account.enabled is True
+
+
+def test_disable(account):
+    """Should disable the account."""
+    account.enabled = True
+    account.disable()
+    assert account.enabled is False
+
+
+def test_log_on_increments(account):
+    """Should increase the log on value by 1."""
+    account.num_logons = 0
+    account.log_on()
+    assert account.num_logons is 1
+
+
+def test_log_off_increments(account):
+    """Should increase the log on value by 1."""
+    account.num_logoffs = 0
+    account.log_off()
+    assert account.num_logoffs is 1
+
+
+def test_account_serialise(account):
     """Test that an account can be serialised. If pydantic throws error then this test fails."""
-    acct = Account(username="Jake", password="JakePass1!", account_type=AccountType.USER)
-    serialised = acct.model_dump_json()
+    serialised = account.model_dump_json()
     print(serialised)
 
 
-def test_account_deserialise():
+def test_account_deserialise(account):
     """Test that an account can be deserialised. The test fails if pydantic throws an error."""
     acct_json = (
         '{"uuid":"dfb2bcaa-d3a1-48fd-af3f-c943354622b4","num_logons":0,"num_logoffs":0,"num_group_changes":0,'
-        '"username":"Jake","password":"JakePass1!","account_type":2,"status":2,"request_manager":null}'
+        '"username":"Jake","password":"totally_hashed_password","account_type":2,"status":2,"request_manager":null}'
     )
-    acct = Account.model_validate_json(acct_json)
+    assert Account.model_validate_json(acct_json)
+
+
+def test_describe_state(account):
+    state = account.describe_state()
+    assert state["num_logons"] is 0
+    assert state["num_logoffs"] is 0
+    assert state["num_group_changes"] is 0
+    assert state["username"] is "Jake"
+    assert state["password"] is "totally_hashed_password"
+    assert state["account_type"] is AccountType.USER.value
+    assert state["enabled"] is True
+
+    account.log_on()
+    state = account.describe_state()
+    assert state["num_logons"] is 1
+    assert state["num_logoffs"] is 0
+    assert state["num_group_changes"] is 0
+    assert state["username"] is "Jake"
+    assert state["password"] is "totally_hashed_password"
+    assert state["account_type"] is AccountType.USER.value
+    assert state["enabled"] is True
+
+    account.log_off()
+    state = account.describe_state()
+    assert state["num_logons"] is 1
+    assert state["num_logoffs"] is 1
+    assert state["num_group_changes"] is 0
+    assert state["username"] is "Jake"
+    assert state["password"] is "totally_hashed_password"
+    assert state["account_type"] is AccountType.USER.value
+    assert state["enabled"] is True
+
+    account.disable()
+    state = account.describe_state()
+    assert state["num_logons"] is 1
+    assert state["num_logoffs"] is 1
+    assert state["num_group_changes"] is 0
+    assert state["username"] is "Jake"
+    assert state["password"] is "totally_hashed_password"
+    assert state["account_type"] is AccountType.USER.value
+    assert state["enabled"] is False
diff --git a/tests/unit_tests/_primaite/_simulator/_network/test_container.py b/tests/unit_tests/_primaite/_simulator/_network/test_container.py
index 66bd59a9..92b3a91b 100644
--- a/tests/unit_tests/_primaite/_simulator/_network/test_container.py
+++ b/tests/unit_tests/_primaite/_simulator/_network/test_container.py
@@ -3,6 +3,64 @@ import json
 import pytest
 
 from primaite.simulator.network.container import Network
+from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
+from primaite.simulator.network.hardware.nodes.computer import Computer
+from primaite.simulator.system.applications.database_client import DatabaseClient
+from primaite.simulator.system.services.database.database_service import DatabaseService
+
+
+@pytest.fixture(scope="function")
+def network(example_network) -> Network:
+    assert len(example_network.routers) is 1
+    assert len(example_network.switches) is 2
+    assert len(example_network.computers) is 2
+    assert len(example_network.servers) is 2
+
+    example_network.set_original_state()
+
+    return example_network
+
+
+def test_describe_state(example_network):
+    """Test that describe state works."""
+    state = example_network.describe_state()
+
+    assert len(state["nodes"]) is 7
+    assert len(state["links"]) is 6
+
+
+def test_reset_network(example_network):
+    """
+    Test that the network is properly reset.
+
+    TODO: make sure that once implemented - any installed/uninstalled services, processes, apps,
+    etc are also removed/reinstalled
+
+    """
+    state_before = example_network.describe_state()
+
+    client_1: Computer = example_network.get_node_by_hostname("client_1")
+    server_1: Computer = example_network.get_node_by_hostname("server_1")
+
+    assert client_1.operating_state is NodeOperatingState.ON
+    assert server_1.operating_state is NodeOperatingState.ON
+
+    client_1.power_off()
+    assert client_1.operating_state is NodeOperatingState.SHUTTING_DOWN
+
+    server_1.power_off()
+    assert server_1.operating_state is NodeOperatingState.SHUTTING_DOWN
+
+    assert example_network.describe_state() is not state_before
+
+    example_network.reset_component_for_episode(episode=1)
+
+    assert client_1.operating_state is NodeOperatingState.ON
+    assert server_1.operating_state is NodeOperatingState.ON
+
+    assert json.dumps(example_network.describe_state(), sort_keys=True, indent=2) == json.dumps(
+        state_before, sort_keys=True, indent=2
+    )
 
 
 def test_creating_container():
@@ -10,11 +68,3 @@ def test_creating_container():
     net = Network()
     assert net.nodes == {}
     assert net.links == {}
-
-
-@pytest.mark.skip(reason="Skipping until we tackle serialisation")
-def test_describe_state():
-    """Check that we can describe network state without raising errors, and that the result is JSON serialisable."""
-    net = Network()
-    state = net.describe_state()
-    json.dumps(state)  # if this function call raises an error, the test fails, state was not JSON-serialisable
diff --git a/tests/unit_tests/_primaite/_simulator/_system/_applications/test_database_client.py b/tests/unit_tests/_primaite/_simulator/_system/_applications/test_database_client.py
new file mode 100644
index 00000000..59d44561
--- /dev/null
+++ b/tests/unit_tests/_primaite/_simulator/_system/_applications/test_database_client.py
@@ -0,0 +1,122 @@
+from ipaddress import IPv4Address
+from typing import Tuple, Union
+
+import pytest
+
+from primaite.simulator.network.hardware.node_operating_state import NodeOperatingState
+from primaite.simulator.network.hardware.nodes.computer import Computer
+from primaite.simulator.system.applications.application import ApplicationOperatingState
+from primaite.simulator.system.applications.database_client import DatabaseClient
+
+
+@pytest.fixture(scope="function")
+def database_client_on_computer() -> Tuple[DatabaseClient, Computer]:
+    computer = Computer(
+        hostname="db_node", ip_address="192.168.0.1", subnet_mask="255.255.255.0", operating_state=NodeOperatingState.ON
+    )
+    computer.software_manager.install(DatabaseClient)
+
+    database_client: DatabaseClient = computer.software_manager.software.get("DatabaseClient")
+    database_client.configure(server_ip_address=IPv4Address("192.168.0.1"))
+    database_client.run()
+    return database_client, computer
+
+
+def test_creation(database_client_on_computer):
+    database_client, computer = database_client_on_computer
+    database_client.describe_state()
+
+
+def test_connect_when_client_is_closed(database_client_on_computer):
+    """Database client should not connect when it is not running."""
+    database_client, computer = database_client_on_computer
+
+    database_client.close()
+    assert database_client.operating_state is ApplicationOperatingState.CLOSED
+
+    assert database_client.connect() is False
+
+
+def test_connect_to_database_fails_on_reattempt(database_client_on_computer):
+    """Database client should return False when the attempt to connect fails."""
+    database_client, computer = database_client_on_computer
+
+    database_client.connected = False
+    assert database_client._connect(server_ip_address=IPv4Address("192.168.0.1"), is_reattempt=True) is False
+
+
+def test_disconnect_when_client_is_closed(database_client_on_computer):
+    """Database client disconnect should not do anything when it is not running."""
+    database_client, computer = database_client_on_computer
+
+    database_client.connected = True
+    assert database_client.server_ip_address is not None
+
+    database_client.close()
+    assert database_client.operating_state is ApplicationOperatingState.CLOSED
+
+    database_client.disconnect()
+
+    assert database_client.connected is True
+    assert database_client.server_ip_address is not None
+
+
+def test_disconnect(database_client_on_computer):
+    """Database client should set connected to False and remove the database server ip address."""
+    database_client, computer = database_client_on_computer
+
+    database_client.connected = True
+
+    assert database_client.operating_state is ApplicationOperatingState.RUNNING
+    assert database_client.server_ip_address is not None
+
+    database_client.disconnect()
+
+    assert database_client.connected is False
+    assert database_client.server_ip_address is None
+
+
+def test_query_when_client_is_closed(database_client_on_computer):
+    """Database client should return False when it is not running."""
+    database_client, computer = database_client_on_computer
+
+    database_client.close()
+    assert database_client.operating_state is ApplicationOperatingState.CLOSED
+
+    assert database_client.query(sql="test") is False
+
+
+def test_query_failed_reattempt(database_client_on_computer):
+    """Database client query should return False if the reattempt fails."""
+    database_client, computer = database_client_on_computer
+
+    def return_false():
+        return False
+
+    database_client.connect = return_false
+
+    database_client.connected = False
+    assert database_client.query(sql="test", is_reattempt=True) is False
+
+
+def test_query_fail_to_connect(database_client_on_computer):
+    """Database client query should return False if the connect attempt fails."""
+    database_client, computer = database_client_on_computer
+
+    def return_false():
+        return False
+
+    database_client.connect = return_false
+    database_client.connected = False
+
+    assert database_client.query(sql="test") is False
+
+
+def test_client_receives_response_when_closed(database_client_on_computer):
+    """Database client receive should return False when it is closed."""
+    database_client, computer = database_client_on_computer
+
+    database_client.close()
+    assert database_client.operating_state is ApplicationOperatingState.CLOSED
+
+    database_client.receive(payload={}, session_id="")
diff --git a/tests/unit_tests/_primaite/_simulator/_system/_applications/test_web_browser.py b/tests/unit_tests/_primaite/_simulator/_system/_applications/test_web_browser.py
index 83426409..dc8f7419 100644
--- a/tests/unit_tests/_primaite/_simulator/_system/_applications/test_web_browser.py
+++ b/tests/unit_tests/_primaite/_simulator/_system/_applications/test_web_browser.py
@@ -21,7 +21,7 @@ def web_browser() -> WebBrowser:
         operating_state=NodeOperatingState.ON,
     )
     # Web Browser should be pre-installed in computer
-    web_browser: WebBrowser = computer.software_manager.software["WebBrowser"]
+    web_browser: WebBrowser = computer.software_manager.software.get("WebBrowser")
     web_browser.run()
     assert web_browser.operating_state is ApplicationOperatingState.RUNNING
     return web_browser
@@ -36,7 +36,7 @@ def test_create_web_client():
         operating_state=NodeOperatingState.ON,
     )
     # Web Browser should be pre-installed in computer
-    web_browser: WebBrowser = computer.software_manager.software["WebBrowser"]
+    web_browser: WebBrowser = computer.software_manager.software.get("WebBrowser")
     assert web_browser.name is "WebBrowser"
     assert web_browser.port is Port.HTTP
     assert web_browser.protocol is IPProtocol.TCP
diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/_red_services/test_data_manipulation_bot.py b/tests/unit_tests/_primaite/_simulator/_system/_services/_red_services/test_data_manipulation_bot.py
index 3b1e4aa4..2c4826bf 100644
--- a/tests/unit_tests/_primaite/_simulator/_system/_services/_red_services/test_data_manipulation_bot.py
+++ b/tests/unit_tests/_primaite/_simulator/_system/_services/_red_services/test_data_manipulation_bot.py
@@ -19,11 +19,11 @@ def dm_client() -> Node:
 
 @pytest.fixture
 def dm_bot(dm_client) -> DataManipulationBot:
-    return dm_client.software_manager.software["DataManipulationBot"]
+    return dm_client.software_manager.software.get("DataManipulationBot")
 
 
 def test_create_dm_bot(dm_client):
-    data_manipulation_bot: DataManipulationBot = dm_client.software_manager.software["DataManipulationBot"]
+    data_manipulation_bot: DataManipulationBot = dm_client.software_manager.software.get("DataManipulationBot")
 
     assert data_manipulation_bot.name == "DataManipulationBot"
     assert data_manipulation_bot.port == Port.POSTGRES_SERVER
diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_database.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_database.py
index 7662fbff..4d96b584 100644
--- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_database.py
+++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_database.py
@@ -8,7 +8,7 @@ from primaite.simulator.system.services.database.database_service import Databas
 def database_server() -> Node:
     node = Node(hostname="db_node")
     node.software_manager.install(DatabaseService)
-    node.software_manager.software["DatabaseService"].start()
+    node.software_manager.software.get("DatabaseService").start()
     return node
 
 
diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns_client.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns_client.py
index 71517855..2bcb512d 100644
--- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns_client.py
+++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns_client.py
@@ -26,14 +26,14 @@ def dns_client() -> Node:
 
 def test_create_dns_client(dns_client):
     assert dns_client is not None
-    dns_client_service: DNSClient = dns_client.software_manager.software["DNSClient"]
+    dns_client_service: DNSClient = dns_client.software_manager.software.get("DNSClient")
     assert dns_client_service.name is "DNSClient"
     assert dns_client_service.port is Port.DNS
     assert dns_client_service.protocol is IPProtocol.TCP
 
 
 def test_dns_client_add_domain_to_cache_when_not_running(dns_client):
-    dns_client_service: DNSClient = dns_client.software_manager.software["DNSClient"]
+    dns_client_service: DNSClient = dns_client.software_manager.software.get("DNSClient")
     assert dns_client.operating_state is NodeOperatingState.OFF
     assert dns_client_service.operating_state is ServiceOperatingState.STOPPED
 
@@ -46,7 +46,7 @@ def test_dns_client_add_domain_to_cache_when_not_running(dns_client):
 
 def test_dns_client_check_domain_exists_when_not_running(dns_client):
     dns_client.operating_state = NodeOperatingState.ON
-    dns_client_service: DNSClient = dns_client.software_manager.software["DNSClient"]
+    dns_client_service: DNSClient = dns_client.software_manager.software.get("DNSClient")
     dns_client_service.start()
 
     assert dns_client.operating_state is NodeOperatingState.ON
@@ -73,7 +73,7 @@ def test_dns_client_check_domain_exists_when_not_running(dns_client):
 def test_dns_client_check_domain_in_cache(dns_client):
     """Test to make sure that the check_domain_in_cache returns the correct values."""
     dns_client.operating_state = NodeOperatingState.ON
-    dns_client_service: DNSClient = dns_client.software_manager.software["DNSClient"]
+    dns_client_service: DNSClient = dns_client.software_manager.software.get("DNSClient")
     dns_client_service.start()
 
     # add a domain to the dns client cache
@@ -85,7 +85,7 @@ def test_dns_client_check_domain_in_cache(dns_client):
 
 def test_dns_client_receive(dns_client):
     """Test to make sure the DNS Client knows how to deal with request responses."""
-    dns_client_service: DNSClient = dns_client.software_manager.software["DNSClient"]
+    dns_client_service: DNSClient = dns_client.software_manager.software.get("DNSClient")
 
     dns_client_service.receive(
         payload=DNSPacket(
@@ -99,6 +99,6 @@ def test_dns_client_receive(dns_client):
 
 
 def test_dns_client_receive_non_dns_payload(dns_client):
-    dns_client_service: DNSClient = dns_client.software_manager.software["DNSClient"]
+    dns_client_service: DNSClient = dns_client.software_manager.software.get("DNSClient")
 
     assert dns_client_service.receive(payload=None) is False
diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns_server.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns_server.py
index 5b65dfc2..eb042c92 100644
--- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns_server.py
+++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_dns_server.py
@@ -26,7 +26,7 @@ def dns_server() -> Node:
 
 def test_create_dns_server(dns_server):
     assert dns_server is not None
-    dns_server_service: DNSServer = dns_server.software_manager.software["DNSServer"]
+    dns_server_service: DNSServer = dns_server.software_manager.software.get("DNSServer")
     assert dns_server_service.name is "DNSServer"
     assert dns_server_service.port is Port.DNS
     assert dns_server_service.protocol is IPProtocol.TCP
@@ -34,7 +34,7 @@ def test_create_dns_server(dns_server):
 
 def test_dns_server_domain_name_registration(dns_server):
     """Test to check if the domain name registration works."""
-    dns_server_service: DNSServer = dns_server.software_manager.software["DNSServer"]
+    dns_server_service: DNSServer = dns_server.software_manager.software.get("DNSServer")
 
     # register the web server in the domain controller
     dns_server_service.dns_register(domain_name="real-domain.com", domain_ip_address=IPv4Address("192.168.1.12"))
@@ -46,7 +46,7 @@ def test_dns_server_domain_name_registration(dns_server):
 
 def test_dns_server_receive(dns_server):
     """Test to make sure that the DNS Server correctly responds to a DNS Client request."""
-    dns_server_service: DNSServer = dns_server.software_manager.software["DNSServer"]
+    dns_server_service: DNSServer = dns_server.software_manager.software.get("DNSServer")
 
     # register the web server in the domain controller
     dns_server_service.dns_register(domain_name="real-domain.com", domain_ip_address=IPv4Address("192.168.1.12"))
diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp_client.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp_client.py
index 1d7355a2..134f82bd 100644
--- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp_client.py
+++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp_client.py
@@ -26,7 +26,7 @@ def ftp_client() -> Node:
 
 def test_create_ftp_client(ftp_client):
     assert ftp_client is not None
-    ftp_client_service: FTPClient = ftp_client.software_manager.software["FTPClient"]
+    ftp_client_service: FTPClient = ftp_client.software_manager.software.get("FTPClient")
     assert ftp_client_service.name is "FTPClient"
     assert ftp_client_service.port is Port.FTP
     assert ftp_client_service.protocol is IPProtocol.TCP
@@ -47,7 +47,7 @@ def test_ftp_client_store_file(ftp_client):
         status_code=FTPStatusCode.OK,
     )
 
-    ftp_client_service: FTPClient = ftp_client.software_manager.software["FTPClient"]
+    ftp_client_service: FTPClient = ftp_client.software_manager.software.get("FTPClient")
     ftp_client_service.receive(response)
 
     assert ftp_client.file_system.get_file(folder_name="downloads", file_name="file.txt")
@@ -61,7 +61,7 @@ def test_ftp_should_not_process_commands_if_service_not_running(ftp_client):
         status_code=FTPStatusCode.OK,
     )
 
-    ftp_client_service: FTPClient = ftp_client.software_manager.software["FTPClient"]
+    ftp_client_service: FTPClient = ftp_client.software_manager.software.get("FTPClient")
     ftp_client_service.stop()
     assert ftp_client_service.operating_state is ServiceOperatingState.STOPPED
     assert ftp_client_service._process_ftp_command(payload=payload).status_code is FTPStatusCode.ERROR
@@ -71,7 +71,7 @@ def test_ftp_tries_to_senf_file__that_does_not_exist(ftp_client):
     """Method send_file should return false if no file to send."""
     assert ftp_client.file_system.get_file(folder_name="root", file_name="test.txt") is None
 
-    ftp_client_service: FTPClient = ftp_client.software_manager.software["FTPClient"]
+    ftp_client_service: FTPClient = ftp_client.software_manager.software.get("FTPClient")
     assert ftp_client_service.operating_state is ServiceOperatingState.RUNNING
     assert (
         ftp_client_service.send_file(
@@ -87,7 +87,7 @@ def test_ftp_tries_to_senf_file__that_does_not_exist(ftp_client):
 
 def test_offline_ftp_client_receives_request(ftp_client):
     """Receive should return false if the node the ftp client is installed on is offline."""
-    ftp_client_service: FTPClient = ftp_client.software_manager.software["FTPClient"]
+    ftp_client_service: FTPClient = ftp_client.software_manager.software.get("FTPClient")
     ftp_client.power_off()
 
     for i in range(ftp_client.shut_down_duration + 1):
@@ -107,7 +107,7 @@ def test_offline_ftp_client_receives_request(ftp_client):
 
 def test_receive_should_fail_if_payload_is_not_ftp(ftp_client):
     """Receive should return false if the node the ftp client is installed on is not an FTPPacket."""
-    ftp_client_service: FTPClient = ftp_client.software_manager.software["FTPClient"]
+    ftp_client_service: FTPClient = ftp_client.software_manager.software.get("FTPClient")
     assert ftp_client_service.receive(payload=None) is False
 
 
@@ -118,5 +118,5 @@ def test_receive_should_ignore_payload_with_none_status_code(ftp_client):
         ftp_command_args=Port.FTP,
         status_code=None,
     )
-    ftp_client_service: FTPClient = ftp_client.software_manager.software["FTPClient"]
+    ftp_client_service: FTPClient = ftp_client.software_manager.software.get("FTPClient")
     assert ftp_client_service.receive(payload=payload) is False
diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp_server.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp_server.py
index 0c849106..2b26c932 100644
--- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp_server.py
+++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_ftp_server.py
@@ -25,7 +25,7 @@ def ftp_server() -> Node:
 
 def test_create_ftp_server(ftp_server):
     assert ftp_server is not None
-    ftp_server_service: FTPServer = ftp_server.software_manager.software["FTPServer"]
+    ftp_server_service: FTPServer = ftp_server.software_manager.software.get("FTPServer")
     assert ftp_server_service.name is "FTPServer"
     assert ftp_server_service.port is Port.FTP
     assert ftp_server_service.protocol is IPProtocol.TCP
@@ -45,7 +45,7 @@ def test_ftp_server_store_file(ftp_server):
         packet_payload_size=24,
     )
 
-    ftp_server_service: FTPServer = ftp_server.software_manager.software["FTPServer"]
+    ftp_server_service: FTPServer = ftp_server.software_manager.software.get("FTPServer")
     ftp_server_service.receive(response)
 
     assert ftp_server.file_system.get_file(folder_name="downloads", file_name="file.txt")
@@ -59,7 +59,7 @@ def test_ftp_server_should_send_error_if_port_arg_is_invalid(ftp_server):
         packet_payload_size=24,
     )
 
-    ftp_server_service: FTPServer = ftp_server.software_manager.software["FTPServer"]
+    ftp_server_service: FTPServer = ftp_server.software_manager.software.get("FTPServer")
     assert ftp_server_service._process_ftp_command(payload=payload).status_code is FTPStatusCode.ERROR
 
 
@@ -67,7 +67,7 @@ def test_ftp_server_receives_non_ftp_packet(ftp_server):
     """Receive should return false if the service receives a non ftp packet."""
     response: FTPPacket = None
 
-    ftp_server_service: FTPServer = ftp_server.software_manager.software["FTPServer"]
+    ftp_server_service: FTPServer = ftp_server.software_manager.software.get("FTPServer")
     assert ftp_server_service.receive(response) is False
 
 
@@ -83,7 +83,7 @@ def test_offline_ftp_server_receives_request(ftp_server):
         packet_payload_size=24,
     )
 
-    ftp_server_service: FTPServer = ftp_server.software_manager.software["FTPServer"]
+    ftp_server_service: FTPServer = ftp_server.software_manager.software.get("FTPServer")
     ftp_server_service.stop()
     assert ftp_server_service.operating_state is ServiceOperatingState.STOPPED
     assert ftp_server_service.receive(response) is False
diff --git a/tests/unit_tests/_primaite/_simulator/_system/_services/test_web_server.py b/tests/unit_tests/_primaite/_simulator/_system/_services/test_web_server.py
index e6f0b9d9..bbccda27 100644
--- a/tests/unit_tests/_primaite/_simulator/_system/_services/test_web_server.py
+++ b/tests/unit_tests/_primaite/_simulator/_system/_services/test_web_server.py
@@ -18,13 +18,13 @@ def web_server() -> Server:
         hostname="web_server", ip_address="192.168.1.10", subnet_mask="255.255.255.0", default_gateway="192.168.1.1"
     )
     node.software_manager.install(software_class=WebServer)
-    node.software_manager.software["WebServer"].start()
+    node.software_manager.software.get("WebServer").start()
     return node
 
 
 def test_create_web_server(web_server):
     assert web_server is not None
-    web_server_service: WebServer = web_server.software_manager.software["WebServer"]
+    web_server_service: WebServer = web_server.software_manager.software.get("WebServer")
     assert web_server_service.name is "WebServer"
     assert web_server_service.port is Port.HTTP
     assert web_server_service.protocol is IPProtocol.TCP
@@ -33,7 +33,7 @@ def test_create_web_server(web_server):
 def test_handling_get_request_not_found_path(web_server):
     payload = HttpRequestPacket(request_method=HttpRequestMethod.GET, request_url="http://domain.com/fake-path")
 
-    web_server_service: WebServer = web_server.software_manager.software["WebServer"]
+    web_server_service: WebServer = web_server.software_manager.software.get("WebServer")
 
     response: HttpResponsePacket = web_server_service._handle_get_request(payload=payload)
     assert response.status_code == HttpStatusCode.NOT_FOUND
@@ -42,7 +42,7 @@ def test_handling_get_request_not_found_path(web_server):
 def test_handling_get_request_home_page(web_server):
     payload = HttpRequestPacket(request_method=HttpRequestMethod.GET, request_url="http://domain.com/")
 
-    web_server_service: WebServer = web_server.software_manager.software["WebServer"]
+    web_server_service: WebServer = web_server.software_manager.software.get("WebServer")
 
     response: HttpResponsePacket = web_server_service._handle_get_request(payload=payload)
     assert response.status_code == HttpStatusCode.OK
@@ -51,7 +51,7 @@ def test_handling_get_request_home_page(web_server):
 def test_process_http_request_get(web_server):
     payload = HttpRequestPacket(request_method=HttpRequestMethod.GET, request_url="http://domain.com/")
 
-    web_server_service: WebServer = web_server.software_manager.software["WebServer"]
+    web_server_service: WebServer = web_server.software_manager.software.get("WebServer")
 
     assert web_server_service._process_http_request(payload=payload) is True
 
@@ -59,6 +59,6 @@ def test_process_http_request_get(web_server):
 def test_process_http_request_method_not_allowed(web_server):
     payload = HttpRequestPacket(request_method=HttpRequestMethod.DELETE, request_url="http://domain.com/")
 
-    web_server_service: WebServer = web_server.software_manager.software["WebServer"]
+    web_server_service: WebServer = web_server.software_manager.software.get("WebServer")
 
     assert web_server_service._process_http_request(payload=payload) is False