diff --git a/.azure/azure-ci-build-pipeline.yaml b/.azure/azure-ci-build-pipeline.yaml index 0bb03594..9070270a 100644 --- a/.azure/azure-ci-build-pipeline.yaml +++ b/.azure/azure-ci-build-pipeline.yaml @@ -86,5 +86,5 @@ stages: displayName: 'Perform PrimAITE Setup' - script: | - pytest -n 4 + pytest -n auto displayName: 'Run tests' diff --git a/CHANGELOG.md b/CHANGELOG.md index d66257b5..dd7e3466 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added +- File System - ability to emulate a node's file system during a simulation + ## [2.0.0] - 2023-07-26 ### Added diff --git a/pyproject.toml b/pyproject.toml index 4982dfd1..74de37df 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,7 +37,7 @@ dependencies = [ "stable-baselines3==1.6.2", "tensorflow==2.12.0", "typer[all]==0.9.0", - "pydantic" + "pydantic==2.1.1" ] [tool.setuptools.dynamic] diff --git a/src/primaite/simulator/core.py b/src/primaite/simulator/core.py index c3130116..5f8ad57c 100644 --- a/src/primaite/simulator/core.py +++ b/src/primaite/simulator/core.py @@ -1,12 +1,21 @@ """Core of the PrimAITE Simulator.""" from abc import abstractmethod from typing import Callable, Dict, List +from uuid import uuid4 from pydantic import BaseModel class SimComponent(BaseModel): - """Extension of pydantic BaseModel with additional methods that must be defined by all classes in the simulator.""" + """Extension of pydantic BaseModel with additional methods that must be defined by all classes in the simulator.""" + + uuid: str + """The component UUID.""" + + def __init__(self, **kwargs): + if not kwargs.get("uuid"): + kwargs["uuid"] = str(uuid4()) + super().__init__(**kwargs) @abstractmethod def describe_state(self) -> Dict: diff --git a/src/primaite/simulator/file_system/__init__.py b/src/primaite/simulator/file_system/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/primaite/simulator/file_system/file_system.py b/src/primaite/simulator/file_system/file_system.py new file mode 100644 index 00000000..3290570e --- /dev/null +++ b/src/primaite/simulator/file_system/file_system.py @@ -0,0 +1,220 @@ +from random import choice +from typing import Dict, Optional + +from primaite import getLogger +from primaite.simulator.core import SimComponent +from primaite.simulator.file_system.file_system_file import FileSystemFile +from primaite.simulator.file_system.file_system_file_type import FileSystemFileType +from primaite.simulator.file_system.file_system_folder import FileSystemFolder + +_LOGGER = getLogger(__name__) + + +class FileSystem(SimComponent): + """Class that contains all the simulation File System.""" + + folders: Dict = {} + """List containing all the folders in the file system.""" + + def describe_state(self) -> Dict: + """ + Get the current state of the FileSystem as a dict. + + :return: A dict containing the current state of the FileSystemFile. + """ + pass + + def get_folders(self) -> Dict: + """Returns the list of folders.""" + return self.folders + + def create_file( + self, + file_name: str, + size: Optional[float] = None, + file_type: Optional[FileSystemFileType] = None, + folder: Optional[FileSystemFolder] = None, + folder_uuid: Optional[str] = None, + ) -> FileSystemFile: + """ + Creates a FileSystemFile and adds it to the list of files. + + If no size or file_type are provided, one will be chosen randomly. + If no folder_uuid or folder is provided, a new folder will be created. + + :param: file_name: The file name + :type: file_name: str + + :param: size: The size the file takes on disk. + :type: size: Optional[float] + + :param: file_type: The type of the file + :type: Optional[FileSystemFileType] + + :param: folder: The folder to add the file to + :type: folder: Optional[FileSystemFolder] + + :param: folder_uuid: The uuid of the folder to add the file to + :type: folder_uuid: Optional[str] + """ + file = None + folder = None + + if file_type is None: + file_type = self.get_random_file_type() + + # if no folder uuid provided, create a folder and add file to it + if folder_uuid is not None: + # otherwise check for existence and add file + folder = self.get_folder_by_id(folder_uuid) + + if folder is not None: + file = FileSystemFile(name=file_name, size=size, file_type=file_type) + folder.add_file(file=file) + else: + # check if a "root" folder exists + folder = self.get_folder_by_name("root") + if folder is None: + # create a root folder + folder = FileSystemFolder(name="root") + + # add file to root folder + file = FileSystemFile(name=file_name, size=size, file_type=file_type) + folder.add_file(file) + self.folders[folder.uuid] = folder + return file + + def create_folder( + self, + folder_name: str, + ) -> FileSystemFolder: + """ + Creates a FileSystemFolder and adds it to the list of folders. + + :param: folder_name: The name of the folder + :type: folder_name: str + """ + folder = FileSystemFolder(name=folder_name) + self.folders[folder.uuid] = folder + return folder + + def delete_file(self, file: Optional[FileSystemFile] = None): + """ + Deletes a file and removes it from the files list. + + :param file: The file to delete + :type file: Optional[FileSystemFile] + """ + # iterate through folders to delete the item with the matching uuid + for key in self.folders: + self.get_folder_by_id(key).remove_file(file) + + def delete_folder(self, folder: FileSystemFolder): + """ + Deletes a folder, removes it from the folders list and removes any child folders and files. + + :param folder: The folder to remove + :type folder: FileSystemFolder + """ + if folder is None or not isinstance(folder, FileSystemFolder): + raise Exception(f"Invalid folder: {folder}") + + if self.folders.get(folder.uuid): + del self.folders[folder.uuid] + else: + _LOGGER.debug(f"File with UUID {folder.uuid} was not found.") + + def move_file(self, file: FileSystemFile, src_folder: FileSystemFolder, target_folder: FileSystemFolder): + """ + Moves a file from one folder to another. + + can provide + + :param: file: The file to move + :type: file: FileSystemFile + + :param: src_folder: The folder where the file is located + :type: FileSystemFolder + + :param: target_folder: The folder where the file should be moved to + :type: FileSystemFolder + """ + # check that the folders exist + if src_folder is None: + raise Exception("Source folder not provided") + + if target_folder is None: + raise Exception("Target folder not provided") + + if file is None: + raise Exception("File to be moved is None") + + # remove file from src + src_folder.remove_file(file) + + # add file to target + target_folder.add_file(file) + + def copy_file(self, file: FileSystemFile, src_folder: FileSystemFolder, target_folder: FileSystemFolder): + """ + Copies a file from one folder to another. + + can provide + + :param: file: The file to move + :type: file: FileSystemFile + + :param: src_folder: The folder where the file is located + :type: FileSystemFolder + + :param: target_folder: The folder where the file should be moved to + :type: FileSystemFolder + """ + if src_folder is None: + raise Exception("Source folder not provided") + + if target_folder is None: + raise Exception("Target folder not provided") + + if file is None: + raise Exception("File to be moved is None") + + # add file to target + target_folder.add_file(file) + + def get_file_by_id(self, file_id: str) -> FileSystemFile: + """Checks if the file exists in any file system folders.""" + for key in self.folders: + file = self.folders[key].get_file_by_id(file_id=file_id) + if file is not None: + return file + + def get_folder_by_name(self, folder_name: str) -> FileSystemFolder: + """ + Returns a the first folder with a matching name. + + :return: Returns the first FileSydtemFolder with a matching name + """ + matching_folder = None + for key in self.folders: + if self.folders[key].name == folder_name: + matching_folder = self.folders[key] + break + return matching_folder + + def get_folder_by_id(self, folder_id: str) -> FileSystemFolder: + """ + Checks if the folder exists. + + :param: folder_id: The id of the folder to find + :type: folder_id: str + """ + return self.folders[folder_id] + + def get_random_file_type(self) -> FileSystemFileType: + """ + Returns a random FileSystemFileTypeEnum. + + :return: A random file type Enum + """ + return choice(list(FileSystemFileType)) diff --git a/src/primaite/simulator/file_system/file_system_file.py b/src/primaite/simulator/file_system/file_system_file.py new file mode 100644 index 00000000..f9fc2e1f --- /dev/null +++ b/src/primaite/simulator/file_system/file_system_file.py @@ -0,0 +1,45 @@ +from random import choice +from typing import Dict + +from primaite.simulator.file_system.file_system_file_type import file_type_sizes_KB, FileSystemFileType +from primaite.simulator.file_system.file_system_item_abc import FileSystemItem + + +class FileSystemFile(FileSystemItem): + """Class that represents a file in the simulation.""" + + file_type: FileSystemFileType = None + """The type of the FileSystemFile""" + + def __init__(self, **kwargs): + """ + Initialise FileSystemFile class. + + :param name: The name of the file. + :type name: str + + :param file_type: The FileSystemFileType of the file + :type file_type: Optional[FileSystemFileType] + + :param size: The size of the FileSystemItem + :type size: Optional[float] + """ + # set random file type if none provided + + # set random file type if none provided + if kwargs.get("file_type") is None: + kwargs["file_type"] = choice(list(FileSystemFileType)) + + # set random file size if none provided + if kwargs.get("size") is None: + kwargs["size"] = file_type_sizes_KB[kwargs["file_type"]] + + super().__init__(**kwargs) + + def describe_state(self) -> Dict: + """ + Get the current state of the FileSystemFile as a dict. + + :return: A dict containing the current state of the FileSystemFile. + """ + pass diff --git a/src/primaite/simulator/file_system/file_system_file_type.py b/src/primaite/simulator/file_system/file_system_file_type.py new file mode 100644 index 00000000..7e2d8706 --- /dev/null +++ b/src/primaite/simulator/file_system/file_system_file_type.py @@ -0,0 +1,124 @@ +from enum import Enum + + +class FileSystemFileType(str, Enum): + """An enumeration of common file types.""" + + UNKNOWN = 0 + "Unknown file type." + + # Text formats + TXT = 1 + "Plain text file." + DOC = 2 + "Microsoft Word document (.doc)" + DOCX = 3 + "Microsoft Word document (.docx)" + PDF = 4 + "Portable Document Format." + HTML = 5 + "HyperText Markup Language file." + XML = 6 + "Extensible Markup Language file." + CSV = 7 + "Comma-Separated Values file." + + # Spreadsheet formats + XLS = 8 + "Microsoft Excel file (.xls)" + XLSX = 9 + "Microsoft Excel file (.xlsx)" + + # Image formats + JPEG = 10 + "JPEG image file." + PNG = 11 + "PNG image file." + GIF = 12 + "GIF image file." + BMP = 13 + "Bitmap image file." + + # Audio formats + MP3 = 14 + "MP3 audio file." + WAV = 15 + "WAV audio file." + + # Video formats + MP4 = 16 + "MP4 video file." + AVI = 17 + "AVI video file." + MKV = 18 + "MKV video file." + FLV = 19 + "FLV video file." + + # Presentation formats + PPT = 20 + "Microsoft PowerPoint file (.ppt)" + PPTX = 21 + "Microsoft PowerPoint file (.pptx)" + + # Web formats + JS = 22 + "JavaScript file." + CSS = 23 + "Cascading Style Sheets file." + + # Programming languages + PY = 24 + "Python script file." + C = 25 + "C source code file." + CPP = 26 + "C++ source code file." + JAVA = 27 + "Java source code file." + + # Compressed file types + RAR = 28 + "RAR archive file." + ZIP = 29 + "ZIP archive file." + TAR = 30 + "TAR archive file." + GZ = 31 + "Gzip compressed file." + + +file_type_sizes_KB = { + FileSystemFileType.UNKNOWN: 0, + FileSystemFileType.TXT: 4, + FileSystemFileType.DOC: 50, + FileSystemFileType.DOCX: 30, + FileSystemFileType.PDF: 100, + FileSystemFileType.HTML: 15, + FileSystemFileType.XML: 10, + FileSystemFileType.CSV: 15, + FileSystemFileType.XLS: 100, + FileSystemFileType.XLSX: 25, + FileSystemFileType.JPEG: 100, + FileSystemFileType.PNG: 40, + FileSystemFileType.GIF: 30, + FileSystemFileType.BMP: 300, + FileSystemFileType.MP3: 5000, + FileSystemFileType.WAV: 25000, + FileSystemFileType.MP4: 25000, + FileSystemFileType.AVI: 50000, + FileSystemFileType.MKV: 50000, + FileSystemFileType.FLV: 15000, + FileSystemFileType.PPT: 200, + FileSystemFileType.PPTX: 100, + FileSystemFileType.JS: 10, + FileSystemFileType.CSS: 5, + FileSystemFileType.PY: 5, + FileSystemFileType.C: 5, + FileSystemFileType.CPP: 10, + FileSystemFileType.JAVA: 10, + FileSystemFileType.RAR: 1000, + FileSystemFileType.ZIP: 1000, + FileSystemFileType.TAR: 1000, + FileSystemFileType.GZ: 800, +} diff --git a/src/primaite/simulator/file_system/file_system_folder.py b/src/primaite/simulator/file_system/file_system_folder.py new file mode 100644 index 00000000..79e19189 --- /dev/null +++ b/src/primaite/simulator/file_system/file_system_folder.py @@ -0,0 +1,69 @@ +from typing import Dict, Optional + +from primaite import getLogger +from primaite.simulator.file_system.file_system_file import FileSystemFile +from primaite.simulator.file_system.file_system_item_abc import FileSystemItem + +_LOGGER = getLogger(__name__) + + +class FileSystemFolder(FileSystemItem): + """Simulation FileSystemFolder.""" + + files: Dict = {} + """List of files stored in the folder.""" + + is_quarantined: bool = False + """Flag that marks the folder as quarantined if true.""" + + def get_file_by_id(self, file_id: str) -> FileSystemFile: + """Return a FileSystemFile with the matching id.""" + return self.files.get(file_id) + + def add_file(self, file: FileSystemFile): + """Adds a file to the folder list.""" + if file is None or not isinstance(file, FileSystemFile): + raise Exception(f"Invalid file: {file}") + + # add to list + self.files[file.uuid] = file + self.size += file.size + + def remove_file(self, file: Optional[FileSystemFile]): + """ + Removes a file from the folder list. + + The method can take a FileSystemFile object or a file id. + + :param: file: The file to remove + :type: Optional[FileSystemFile] + """ + if file is None or not isinstance(file, FileSystemFile): + raise Exception(f"Invalid file: {file}") + + if self.files.get(file.uuid): + del self.files[file.uuid] + + self.size -= file.size + else: + _LOGGER.debug(f"File with UUID {file.uuid} was not found.") + + def quarantine(self): + """Quarantines the File System Folder.""" + self.is_quarantined = True + + def end_quarantine(self): + """Ends the quarantine of the File System Folder.""" + self.is_quarantined = False + + def quarantine_status(self) -> bool: + """Returns true if the folder is being quarantined.""" + return self.is_quarantined + + def describe_state(self) -> Dict: + """ + Get the current state of the FileSystemFolder as a dict. + + :return: A dict containing the current state of the FileSystemFile. + """ + pass diff --git a/src/primaite/simulator/file_system/file_system_item_abc.py b/src/primaite/simulator/file_system/file_system_item_abc.py new file mode 100644 index 00000000..0594cc35 --- /dev/null +++ b/src/primaite/simulator/file_system/file_system_item_abc.py @@ -0,0 +1,17 @@ +from typing import Dict + +from primaite.simulator.core import SimComponent + + +class FileSystemItem(SimComponent): + """Abstract base class for FileSystemItems used in the file system simulation.""" + + name: str + """The name of the FileSystemItem.""" + + size: float = 0 + """The size the item takes up on disk.""" + + def describe_state(self) -> Dict: + """Returns the state of the FileSystemItem.""" + pass diff --git a/tests/conftest.py b/tests/conftest.py index f40b0b94..f1c05187 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -112,7 +112,7 @@ def temp_session_path() -> Path: session_timestamp = datetime.now() date_dir = session_timestamp.strftime("%Y-%m-%d") session_path = session_timestamp.strftime("%Y-%m-%d_%H-%M-%S") - session_path = Path(tempfile.gettempdir()) / "primaite" / date_dir / session_path + session_path = Path(tempfile.gettempdir()) / "_primaite" / date_dir / session_path session_path.mkdir(exist_ok=True, parents=True) return session_path diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/__init__.py b/tests/unit_tests/_primaite/_simulator/_file_system/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system.py new file mode 100644 index 00000000..5bebf487 --- /dev/null +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system.py @@ -0,0 +1,129 @@ +from primaite.simulator.file_system.file_system import FileSystem +from primaite.simulator.file_system.file_system_folder import FileSystemFolder + + +def test_create_folder_and_file(): + """Test creating a folder and a file.""" + file_system = FileSystem() + folder = file_system.create_folder(folder_name="test_folder") + assert len(file_system.folders) is 1 + + file = file_system.create_file(file_name="test_file", size=10, folder_uuid=folder.uuid) + assert len(file_system.get_folder_by_id(folder.uuid).files) is 1 + + assert file_system.get_file_by_id(file.uuid).name is "test_file" + assert file_system.get_file_by_id(file.uuid).size == 10 + + +def test_create_file(): + """Tests that creating a file without a folder creates a folder and sets that as the file's parent.""" + file_system = FileSystem() + + file = file_system.create_file(file_name="test_file", size=10) + assert len(file_system.folders) is 1 + assert file_system.get_folder_by_name("root").get_file_by_id(file.uuid) is file + + +def test_delete_file(): + """Tests that a file can be deleted.""" + file_system = FileSystem() + + file = file_system.create_file(file_name="test_file", size=10) + assert len(file_system.folders) is 1 + + folder_id = list(file_system.folders.keys())[0] + folder = file_system.get_folder_by_id(folder_id) + assert folder.get_file_by_id(file.uuid) is file + + file_system.delete_file(file=file) + assert len(file_system.folders) is 1 + assert len(file_system.get_folder_by_id(folder.uuid).files) is 0 + + +def test_delete_non_existent_file(): + """Tests deleting a non existent file.""" + file_system = FileSystem() + + file = file_system.create_file(file_name="test_file", size=10) + not_added_file = file_system.create_file(file_name="test_file", size=10) + assert len(file_system.folders) is 1 + + folder_id = list(file_system.folders.keys())[0] + folder = file_system.get_folder_by_id(folder_id) + assert folder.get_file_by_id(file.uuid) is file + + file_system.delete_file(file=not_added_file) + assert len(file_system.folders) is 1 + assert len(file_system.get_folder_by_id(folder.uuid).files) is 1 + + +def test_delete_folder(): + file_system = FileSystem() + folder = file_system.create_folder(folder_name="test_folder") + assert len(file_system.folders) is 1 + + file_system.delete_folder(folder) + assert len(file_system.folders) is 0 + + +def test_deleting_a_non_existent_folder(): + file_system = FileSystem() + folder = file_system.create_folder(folder_name="test_folder") + not_added_folder = FileSystemFolder(name="fake_folder") + assert len(file_system.folders) is 1 + + file_system.delete_folder(not_added_folder) + assert len(file_system.folders) is 1 + + +def test_move_file(): + """Tests the file move function.""" + file_system = FileSystem() + src_folder = file_system.create_folder(folder_name="test_folder_1") + assert len(file_system.folders) is 1 + + target_folder = file_system.create_folder(folder_name="test_folder_2") + assert len(file_system.folders) is 2 + + file = file_system.create_file(file_name="test_file", size=10, folder_uuid=src_folder.uuid) + assert len(file_system.get_folder_by_id(src_folder.uuid).files) is 1 + assert len(file_system.get_folder_by_id(target_folder.uuid).files) is 0 + + file_system.move_file(file=file, src_folder=src_folder, target_folder=target_folder) + + assert len(file_system.get_folder_by_id(src_folder.uuid).files) is 0 + assert len(file_system.get_folder_by_id(target_folder.uuid).files) is 1 + + +def test_copy_file(): + """Tests the file copy function.""" + file_system = FileSystem() + src_folder = file_system.create_folder(folder_name="test_folder_1") + assert len(file_system.folders) is 1 + + target_folder = file_system.create_folder(folder_name="test_folder_2") + assert len(file_system.folders) is 2 + + file = file_system.create_file(file_name="test_file", size=10, folder_uuid=src_folder.uuid) + assert len(file_system.get_folder_by_id(src_folder.uuid).files) is 1 + assert len(file_system.get_folder_by_id(target_folder.uuid).files) is 0 + + file_system.copy_file(file=file, src_folder=src_folder, target_folder=target_folder) + + assert len(file_system.get_folder_by_id(src_folder.uuid).files) is 1 + assert len(file_system.get_folder_by_id(target_folder.uuid).files) is 1 + + +def test_serialisation(): + """Test to check that the object serialisation works correctly.""" + file_system = FileSystem() + folder = file_system.create_folder(folder_name="test_folder") + assert len(file_system.folders) is 1 + + file_system.create_file(file_name="test_file", size=10, folder_uuid=folder.uuid) + assert file_system.get_folder_by_id(folder.uuid) is folder + + serialised_file_sys = file_system.model_dump_json() + deserialised_file_sys = FileSystem.model_validate_json(serialised_file_sys) + + assert file_system.model_dump_json() == deserialised_file_sys.model_dump_json() diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system_file.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system_file.py new file mode 100644 index 00000000..629b9bb9 --- /dev/null +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system_file.py @@ -0,0 +1,23 @@ +from primaite.simulator.file_system.file_system_file import FileSystemFile +from primaite.simulator.file_system.file_system_file_type import FileSystemFileType + + +def test_file_type(): + """Tests tha the FileSystemFile type is set correctly.""" + file = FileSystemFile(name="test", file_type=FileSystemFileType.DOC) + assert file.file_type is FileSystemFileType.DOC + + +def test_get_size(): + """Tests that the file size is being returned properly.""" + file = FileSystemFile(name="test", size=1.5) + assert file.size == 1.5 + + +def test_serialisation(): + """Test to check that the object serialisation works correctly.""" + file = FileSystemFile(name="test", size=1.5, file_type=FileSystemFileType.DOC) + serialised_file = file.model_dump_json() + deserialised_file = FileSystemFile.model_validate_json(serialised_file) + + assert file.model_dump_json() == deserialised_file.model_dump_json() diff --git a/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system_folder.py b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system_folder.py new file mode 100644 index 00000000..1940e886 --- /dev/null +++ b/tests/unit_tests/_primaite/_simulator/_file_system/test_file_system_folder.py @@ -0,0 +1,75 @@ +from primaite.simulator.file_system.file_system_file import FileSystemFile +from primaite.simulator.file_system.file_system_file_type import FileSystemFileType +from primaite.simulator.file_system.file_system_folder import FileSystemFolder + + +def test_adding_removing_file(): + """Test the adding and removing of a file from a folder.""" + folder = FileSystemFolder(name="test") + + file = FileSystemFile(name="test_file", size=10, file_type=FileSystemFileType.DOC) + + folder.add_file(file) + assert folder.size == 10 + assert len(folder.files) is 1 + + folder.remove_file(file) + assert folder.size == 0 + assert len(folder.files) is 0 + + +def test_remove_non_existent_file(): + """Test the removing of a file that does not exist.""" + folder = FileSystemFolder(name="test") + + file = FileSystemFile(name="test_file", size=10, file_type=FileSystemFileType.DOC) + not_added_file = FileSystemFile(name="fake_file", size=10, file_type=FileSystemFileType.DOC) + + folder.add_file(file) + assert folder.size == 10 + assert len(folder.files) is 1 + + folder.remove_file(not_added_file) + assert folder.size == 10 + assert len(folder.files) is 1 + + +def test_get_file_by_id(): + """Test to make sure that the correct file is returned.""" + folder = FileSystemFolder(name="test") + + file = FileSystemFile(name="test_file", size=10, file_type=FileSystemFileType.DOC) + file2 = FileSystemFile(name="test_file_2", size=10, file_type=FileSystemFileType.DOC) + + folder.add_file(file) + folder.add_file(file2) + assert folder.size == 20 + assert len(folder.files) is 2 + + assert folder.get_file_by_id(file_id=file.uuid) is file + + +def test_folder_quarantine_state(): + """Tests the changing of folder quarantine status.""" + folder = FileSystemFolder(name="test") + + assert folder.quarantine_status() is False + + folder.quarantine() + assert folder.quarantine_status() is True + + folder.end_quarantine() + assert folder.quarantine_status() is False + + +def test_serialisation(): + """Test to check that the object serialisation works correctly.""" + folder = FileSystemFolder(name="test") + file = FileSystemFile(name="test_file", size=10, file_type=FileSystemFileType.DOC) + folder.add_file(file) + + serialised_folder = folder.model_dump_json() + + deserialised_folder = FileSystemFolder.model_validate_json(serialised_folder) + + assert folder.model_dump_json() == deserialised_folder.model_dump_json() diff --git a/tests/unit_tests/_primaite/_simulator/test_core.py b/tests/unit_tests/_primaite/_simulator/test_core.py index de0732f9..4e2df757 100644 --- a/tests/unit_tests/_primaite/_simulator/test_core.py +++ b/tests/unit_tests/_primaite/_simulator/test_core.py @@ -42,8 +42,8 @@ class TestIsolatedSimComponent: return {} comp = TestComponent(name="computer", size=(5, 10)) - dump = comp.model_dump() - assert dump == {"name": "computer", "size": (5, 10)} + dump = comp.model_dump_json() + assert comp == TestComponent.model_validate_json(dump) def test_apply_action(self): """Validate that we can override apply_action behaviour and it updates the state of the component."""