#1714: conver file and folder lists to dicts + fixing and adding a few more tests

This commit is contained in:
Czar Echavez
2023-08-07 14:49:59 +01:00
parent b58a3a3e24
commit 554619e4b4
6 changed files with 128 additions and 40 deletions

View File

@@ -1,16 +1,19 @@
from random import choice
from typing import Dict, List, Optional
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: List[FileSystemFolder] = []
folders: Dict = {}
"""List containing all the folders in the file system."""
def describe_state(self) -> Dict:
@@ -21,7 +24,7 @@ class FileSystem(SimComponent):
"""
pass
def get_folders(self) -> List[FileSystemFolder]:
def get_folders(self) -> Dict:
"""Returns the list of folders."""
return self.folders
@@ -48,6 +51,9 @@ class FileSystem(SimComponent):
: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]
"""
@@ -75,16 +81,21 @@ class FileSystem(SimComponent):
# add file to root folder
file = FileSystemFile(item_name=file_name, item_size=file_size, file_type=file_type)
folder.add_file(file)
self.folders.append(folder)
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."""
"""
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(item_name=folder_name)
self.folders.append(folder)
self.folders[folder.uuid] = folder
return folder
def delete_file(self, file: Optional[FileSystemFile] = None):
@@ -93,13 +104,10 @@ class FileSystem(SimComponent):
:param file: The file to delete
:type file: Optional[FileSystemFile]
:param file_id: The UUID of the file item to delete
:type file_id: Optional[str]
"""
# iterate through folders to delete the item with the matching uuid
for folder in self.folders:
folder.remove_file(file=file)
for key in self.folders:
self.get_folder_by_id(key).remove_file(file)
def delete_folder(self, folder: FileSystemFolder):
"""
@@ -108,7 +116,13 @@ class FileSystem(SimComponent):
:param folder: The folder to remove
:type folder: FileSystemFolder
"""
self.folders.remove(folder)
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):
"""
@@ -170,8 +184,8 @@ class FileSystem(SimComponent):
def get_file_by_id(self, file_id: str) -> FileSystemFile:
"""Checks if the file exists in any file system folders."""
for folder in self.folders:
file = folder.get_file(file_id=file_id)
for key in self.folders:
file = self.folders[key].get_file(file_id=file_id)
if file is not None:
return file
@@ -181,12 +195,26 @@ class FileSystem(SimComponent):
:return: Returns the first FileSydtemFolder with a matching name
"""
return next((f for f in self.folders if f.get_folder_name() == folder_name), None)
matching_folder = None
for key in self.folders:
if self.folders[key].get_folder_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."""
return next((f for f in self.folders if f.uuid == folder_id), None)
"""
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."""
"""
Returns a random FileSystemFileTypeEnum.
:return: A random file type Enum
"""
return choice(list(FileSystemFileType))

View File

@@ -1,13 +1,16 @@
from typing import Dict, List, Optional
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: List[FileSystemFile] = []
files: Dict = {}
"""List of files stored in the folder."""
is_quarantined: bool = False
@@ -21,20 +24,23 @@ class FileSystemFolder(FileSystemItem):
"""Returns the item_size of the folder."""
return self.item_size
def get_files(self) -> List[FileSystemFile]:
"""Returns the list of files the folder contains."""
def get_files(self) -> Dict:
"""Returns the files dictionary."""
return self.files
def get_file(self, file_id: str) -> FileSystemFile:
"""Return a FileSystemFile with the matching id."""
return next((f for f in self.files if f.uuid == file_id), None)
return self.files[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}")
self.item_size += file.get_file_size()
# add to list
self.files.append(file)
self.files[file.uuid] = file
def remove_file(self, file: Optional[FileSystemFile]):
"""
@@ -45,10 +51,16 @@ class FileSystemFolder(FileSystemItem):
:param: file: The file to remove
:type: Optional[FileSystemFile]
"""
self.files.remove(file)
if file is None or not isinstance(file, FileSystemFile):
raise Exception(f"Invalid file: {file}")
# remove folder size from folder
self.item_size -= file.get_file_size()
if self.files.get(file.uuid):
del self.files[file.uuid]
# remove folder size from folder
self.item_size -= file.get_file_size()
else:
_LOGGER.debug(f"File with UUID {file.uuid} was not found.")
def quarantine(self):
"""Quarantines the File System Folder."""

View File

@@ -1,4 +1,5 @@
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():
@@ -7,10 +8,11 @@ def test_create_folder_and_file():
folder = file_system.create_folder(folder_name="test_folder")
assert len(file_system.get_folders()) is 1
file_system.create_file(file_name="test_file", file_size=10, folder_uuid=folder.uuid)
assert len(file_system.get_folders()[0].get_files()) is 1
assert file_system.get_folders()[0].get_files()[0].get_file_name() is "test_file"
assert file_system.get_folders()[0].get_files()[0].get_file_size() == 10
file = file_system.create_file(file_name="test_file", file_size=10, folder_uuid=folder.uuid)
assert len(file_system.get_folder_by_id(folder.uuid).get_files()) is 1
assert file_system.get_file_by_id(file.uuid).get_file_name() is "test_file"
assert file_system.get_file_by_id(file.uuid).get_file_size() == 10
def test_create_file():
@@ -19,7 +21,7 @@ def test_create_file():
file = file_system.create_file(file_name="test_file", file_size=10)
assert len(file_system.get_folders()) is 1
assert file_system.get_folders()[0].get_file(file.uuid) is file
assert file_system.get_folder_by_name("root").get_file(file.uuid) is file
def test_delete_file():
@@ -28,11 +30,31 @@ def test_delete_file():
file = file_system.create_file(file_name="test_file", file_size=10)
assert len(file_system.get_folders()) is 1
assert file_system.get_folders()[0].get_file(file.uuid) is file
folder_id = list(file_system.get_folders().keys())[0]
folder = file_system.get_folder_by_id(folder_id)
assert folder.get_file(file.uuid) is file
file_system.delete_file(file=file)
assert len(file_system.get_folders()) is 1
assert len(file_system.get_folders()[0].get_files()) is 0
assert len(file_system.get_folder_by_id(folder.uuid).get_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", file_size=10)
not_added_file = file_system.create_file(file_name="test_file", file_size=10)
assert len(file_system.get_folders()) is 1
folder_id = list(file_system.get_folders().keys())[0]
folder = file_system.get_folder_by_id(folder_id)
assert folder.get_file(file.uuid) is file
file_system.delete_file(file=not_added_file)
assert len(file_system.get_folders()) is 1
assert len(file_system.get_folder_by_id(folder.uuid).get_files()) is 1
def test_delete_folder():
@@ -44,6 +66,16 @@ def test_delete_folder():
assert len(file_system.get_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(item_name="fake_folder")
assert len(file_system.get_folders()) is 1
file_system.delete_folder(not_added_folder)
assert len(file_system.get_folders()) is 1
def test_move_file():
"""Tests the file move function."""
file_system = FileSystem()
@@ -89,9 +121,9 @@ def test_serialisation():
assert len(file_system.get_folders()) is 1
file_system.create_file(file_name="test_file", file_size=10, folder_uuid=folder.uuid)
assert len(file_system.get_folders()[0].get_files()) is 1
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 == deserialised_file_sys
assert file_system.model_dump_json() == deserialised_file_sys.model_dump_json()

View File

@@ -20,4 +20,4 @@ def test_serialisation():
serialised_file = file.model_dump_json()
deserialised_file = FileSystemFile.model_validate_json(serialised_file)
assert file == deserialised_file
assert file.model_dump_json() == deserialised_file.model_dump_json()

View File

@@ -18,6 +18,22 @@ def test_adding_removing_file():
assert len(folder.get_files()) is 0
def test_remove_non_existent_file():
"""Test the removing of a file that does not exist."""
folder = FileSystemFolder(item_name="test")
file = FileSystemFile(item_name="test_file", item_size=10, file_type=FileSystemFileType.DOC)
not_added_file = FileSystemFile(item_name="fake_file", item_size=10, file_type=FileSystemFileType.DOC)
folder.add_file(file)
assert folder.get_folder_size() == 10
assert len(folder.get_files()) is 1
folder.remove_file(not_added_file)
assert folder.get_folder_size() == 10
assert len(folder.get_files()) is 1
def test_get_file_by_id():
"""Test to make sure that the correct file is returned."""
folder = FileSystemFolder(item_name="test")
@@ -56,4 +72,4 @@ def test_serialisation():
deserialised_folder = FileSystemFolder.model_validate_json(serialised_folder)
assert folder == deserialised_folder
assert folder.model_dump_json() == deserialised_folder.model_dump_json()

View File

@@ -42,8 +42,8 @@ class TestIsolatedSimComponent:
return {}
comp = TestComponent(name="computer", size=(5, 10))
dump = comp.model_dump()
assert dump["name"] is "computer"
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."""