Merge remote-tracking branch 'origin/dev' into feature/2084-add-more-tests-to-make-primaite-more-robust

src/primaite/simulator/system/services/ftp/ftp_client.py

@@ -1,6 +1,7 @@
 from ipaddress import IPv4Address
 from typing import Optional
 
+from primaite import getLogger
 from primaite.simulator.file_system.file_system import File
 from primaite.simulator.network.protocols.ftp import FTPCommand, FTPPacket, FTPStatusCode
 from primaite.simulator.network.transmission.network_layer import IPProtocol
@@ -8,6 +9,8 @@ from primaite.simulator.network.transmission.transport_layer import Port
 from primaite.simulator.system.core.software_manager import SoftwareManager
 from primaite.simulator.system.services.ftp.ftp_service import FTPServiceABC
 
+_LOGGER = getLogger(__name__)
+
 
 class FTPClient(FTPServiceABC):
     """
@@ -27,6 +30,18 @@ class FTPClient(FTPServiceABC):
         super().__init__(**kwargs)
         self.start()
 
+    def set_original_state(self):
+        """Sets the original state."""
+        _LOGGER.debug(f"Setting FTPClient original state on node {self.software_manager.node.hostname}")
+        super().set_original_state()
+        vals_to_include = {"connected"}
+        self._original_state.update(self.model_dump(include=vals_to_include))
+
+    def reset_component_for_episode(self, episode: int):
+        """Reset the original state of the SimComponent."""
+        _LOGGER.debug(f"Resetting FTPClient state on node {self.software_manager.node.hostname}")
+        super().reset_component_for_episode(episode)
+
     def _process_ftp_command(self, payload: FTPPacket, session_id: Optional[str] = None, **kwargs) -> FTPPacket:
         """
         Process the command in the FTP Packet.