From 2ecc142c289dac9093445531cbd7cbf147e0c1a1 Mon Sep 17 00:00:00 2001
From: Nick Todd <nick.todd@t-s.qinetiq.com>
Date: Wed, 11 Dec 2024 16:50:43 +0000
Subject: [PATCH] #2888: Changes to Applications and Services previously
 missed.

---
 .../simulator/network/hardware/base.py         | 18 ++++++++++++++++--
 tests/conftest.py                              | 16 +++++++++++++++-
 .../applications/extended_application.py       |  7 +++++++
 .../network/test_broadcast.py                  |  9 ++++++++-
 4 files changed, 46 insertions(+), 4 deletions(-)

diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py
index 51e200e7..02270e38 100644
--- a/src/primaite/simulator/network/hardware/base.py
+++ b/src/primaite/simulator/network/hardware/base.py
@@ -824,7 +824,7 @@ class User(SimComponent):
         return self.model_dump()
 
 
-class UserManager(Service):
+class UserManager(Service, identifier="UserManager"):
     """
     Manages users within the PrimAITE system, handling creation, authentication, and administration.
 
@@ -833,8 +833,15 @@ class UserManager(Service):
     :param disabled_admins: A dictionary of currently disabled admin users by their usernames
     """
 
+    config: "UserManager.ConfigSchema"
+
     users: Dict[str, User] = {}
 
+    class ConfigSchema(Service.ConfigSchema):
+        """ConfigSchema for UserManager."""
+
+        type: str = "USERMANAGER"
+
     def __init__(self, **kwargs):
         """
         Initializes a UserManager instanc.
@@ -1130,13 +1137,15 @@ class RemoteUserSession(UserSession):
         return state
 
 
-class UserSessionManager(Service):
+class UserSessionManager(Service, identifier="UserSessionManager"):
     """
     Manages user sessions on a Node, including local and remote sessions.
 
     This class handles authentication, session management, and session timeouts for users interacting with the Node.
     """
 
+    config: "UserSessionManager.ConfigSchema"
+
     local_session: Optional[UserSession] = None
     """The current local user session, if any."""
 
@@ -1158,6 +1167,11 @@ class UserSessionManager(Service):
     current_timestep: int = 0
     """The current timestep in the simulation."""
 
+    class ConfigSchema(Service.ConfigSchema):
+        """ConfigSchema for UserSessionManager."""
+
+        type: str = "USERSESSIONMANAGER"
+
     def __init__(self, **kwargs):
         """
         Initializes a UserSessionManager instance.
diff --git a/tests/conftest.py b/tests/conftest.py
index 64fe0699..071d7d99 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -37,9 +37,16 @@ ACTION_SPACE_NODE_ACTION_VALUES = 1
 _LOGGER = getLogger(__name__)
 
 
-class DummyService(Service):
+class DummyService(Service, identifier="DummyService"):
     """Test Service class"""
 
+    config: "DummyService.ConfigSchema"
+
+    class ConfigSchema(Service.ConfigSchema):
+        """ConfigSchema for DummyService."""
+
+        type: str = "DUMMYSERVICE"
+
     def describe_state(self) -> Dict:
         return super().describe_state()
 
@@ -56,6 +63,13 @@ class DummyService(Service):
 class DummyApplication(Application, identifier="DummyApplication"):
     """Test Application class"""
 
+    config: "DummyApplication.ConfigSchema"
+
+    class ConfigSchema(Application.ConfigSchema):
+        """ConfigSchema for DummyApplication."""
+
+        type: str = "DUMMYAPPLICATION"
+
     def __init__(self, **kwargs):
         kwargs["name"] = "DummyApplication"
         kwargs["port"] = PORT_LOOKUP["HTTP"]
diff --git a/tests/integration_tests/extensions/applications/extended_application.py b/tests/integration_tests/extensions/applications/extended_application.py
index 70dc7cba..189d7975 100644
--- a/tests/integration_tests/extensions/applications/extended_application.py
+++ b/tests/integration_tests/extensions/applications/extended_application.py
@@ -31,6 +31,8 @@ class ExtendedApplication(Application, identifier="ExtendedApplication"):
     The application requests and loads web pages using its domain name and requesting IP addresses using DNS.
     """
 
+    config: "ExtendedApplication.ConfigSchema"
+
     target_url: Optional[str] = None
 
     domain_name_ip_address: Optional[IPv4Address] = None
@@ -42,6 +44,11 @@ class ExtendedApplication(Application, identifier="ExtendedApplication"):
     history: List["BrowserHistoryItem"] = []
     """Keep a log of visited websites and information about the visit, such as response code."""
 
+    class ConfigSchema(Application.ConfigSchema):
+        """ConfigSchema for ExtendedApplication."""
+
+        type: str = "EXTENDEDAPPLICATION"
+
     def __init__(self, **kwargs):
         kwargs["name"] = "ExtendedApplication"
         kwargs["protocol"] = PROTOCOL_LOOKUP["TCP"]
diff --git a/tests/integration_tests/network/test_broadcast.py b/tests/integration_tests/network/test_broadcast.py
index f07f02e7..675e0f53 100644
--- a/tests/integration_tests/network/test_broadcast.py
+++ b/tests/integration_tests/network/test_broadcast.py
@@ -14,9 +14,16 @@ from primaite.utils.validation.ip_protocol import PROTOCOL_LOOKUP
 from primaite.utils.validation.port import PORT_LOOKUP
 
 
-class BroadcastTestService(Service):
+class BroadcastTestService(Service, identifier="BroadcastTestService"):
     """A service for sending broadcast and unicast messages over a network."""
 
+    config: "BroadcastTestService.ConfigSchema"
+
+    class ConfigSchema(Service.ConfigSchema):
+        """ConfigSchema for BroadcastTestService."""
+
+        type: str = "BROADCASTTESTSERVICE"
+
     def __init__(self, **kwargs):
         # Set default service properties for broadcasting
         kwargs["name"] = "BroadcastService"