#1714: conver file and folder lists to dicts + fixing and adding a few more tests
This commit is contained in:
@@ -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))
|
||||
|
||||
@@ -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."""
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user