#1714: updated file system classes
This commit is contained in:
@@ -1,6 +1,7 @@
|
||||
"""Core of the PrimAITE Simulator."""
|
||||
from abc import abstractmethod
|
||||
from typing import Callable, Dict, List
|
||||
from uuid import uuid4
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
@@ -8,6 +9,14 @@ from pydantic import BaseModel
|
||||
class SimComponent(BaseModel):
|
||||
"""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:
|
||||
"""
|
||||
|
||||
@@ -1,17 +1,17 @@
|
||||
from typing import Dict, List, Union
|
||||
from typing import Dict, List, Optional
|
||||
|
||||
from pydantic import PrivateAttr
|
||||
|
||||
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
|
||||
|
||||
|
||||
class FileSystem(SimComponent):
|
||||
"""Class that contains all the simulation File System."""
|
||||
|
||||
files: List[FileSystemFile]
|
||||
"""List containing all the files in the file system."""
|
||||
|
||||
folders: List[FileSystemFolder]
|
||||
_folders: List[FileSystemFolder] = PrivateAttr([])
|
||||
"""List containing all the folders in the file system."""
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
@@ -22,56 +22,107 @@ class FileSystem(SimComponent):
|
||||
"""
|
||||
pass
|
||||
|
||||
def create_file(self):
|
||||
"""Creates a FileSystemFile and adds it to the list of files."""
|
||||
pass
|
||||
def get_folders(self) -> List[FileSystemFolder]:
|
||||
"""Returns the list of folders."""
|
||||
return self._folders
|
||||
|
||||
def create_folder(self):
|
||||
def create_file(self, file_size: float, folder_uuid: Optional[str] = None) -> FileSystemFile:
|
||||
"""
|
||||
Creates a FileSystemFile and adds it to the list of files.
|
||||
|
||||
:param: folder_uuid: The uuid of the folder to add the file to
|
||||
:type: folder_uuid: str
|
||||
"""
|
||||
file = None
|
||||
# if no folder uuid provided, create a folder and add file to it
|
||||
if folder_uuid is None:
|
||||
folder = FileSystemFolder()
|
||||
|
||||
file = FileSystemFile(item_parent=folder.uuid, file_size=file_size, file_type=FileSystemFileType.TBD)
|
||||
folder.add_file(file)
|
||||
self._folders.append(folder)
|
||||
else:
|
||||
# otherwise check for existence and add file
|
||||
folder = self.get_folder_by_id(folder_uuid)
|
||||
if folder is not None:
|
||||
file = FileSystemFile(file_size=file_size, file_type=FileSystemFileType.TBD)
|
||||
folder.add_file(file=file)
|
||||
return file
|
||||
|
||||
def create_folder(self) -> FileSystemFolder:
|
||||
"""Creates a FileSystemFolder and adds it to the list of folders."""
|
||||
pass
|
||||
folder = FileSystemFolder(item_parent=None)
|
||||
self._folders.append(folder)
|
||||
return folder
|
||||
|
||||
def delete_file(self, file_item: str):
|
||||
def delete_file(self, file_id: str):
|
||||
"""
|
||||
Deletes a file and removes it from the files list.
|
||||
|
||||
:param file_item: The UUID of the file item to delete
|
||||
:type file_item: str
|
||||
:param file_id: The UUID of the file item to delete
|
||||
:type file_id: str
|
||||
"""
|
||||
self.files = list(filter(lambda x: (x.get_item_uuid() != file_item), self.files))
|
||||
# iterate through folders to delete the item with the matching uuid
|
||||
for folder in self._folders:
|
||||
folder.remove_file(file_id)
|
||||
|
||||
def delete_folder(self, file_item: str):
|
||||
def delete_folder(self, folder_id: str):
|
||||
"""
|
||||
Deletes a folder, removes it frdom the folders list and removes any child folders and files.
|
||||
Deletes a folder, removes it from the folders list and removes any child folders and files.
|
||||
|
||||
:param file_item: The UUID of the file item to delete
|
||||
:type file_item: str
|
||||
:param folder_id: The UUID of the file item to delete
|
||||
:type folder_id: str
|
||||
"""
|
||||
self.files = list(filter(lambda x: (x.get_item_parent() != file_item), self.files))
|
||||
self.folders = list(filter(lambda x: (x.get_item_uuid() != file_item), self.folders))
|
||||
self._folders = list(filter(lambda f: (f.uuid != folder_id), self._folders))
|
||||
|
||||
def move_file_item(self, file_item: str, target_directory: str):
|
||||
"""
|
||||
Check to see if the file_item and target_directory exists then moves the item by changing its parent item uuid.
|
||||
def move_file(self, src_folder_id: str, target_folder_id: str, file_id: str):
|
||||
"""Moves a file from one folder to another."""
|
||||
# check that both folders and the file exists
|
||||
src = self.get_folder_by_id(src_folder_id)
|
||||
target = self.get_folder_by_id(target_folder_id)
|
||||
|
||||
:param file_item: The UUID of the file item to move
|
||||
:type file_item: str
|
||||
if src is None:
|
||||
raise Exception(f"src folder with UUID {src_folder_id} could not be found")
|
||||
|
||||
:param target_directory: The UUID of the directory the item should be moved into
|
||||
:type target_directory: str
|
||||
"""
|
||||
item = self._file_item_exists(file_item)
|
||||
if item and any(f for f in self.folders if f.get_item_uuid() == target_directory):
|
||||
item.move(target_directory)
|
||||
if target is None:
|
||||
raise Exception(f"src folder with UUID {target_folder_id} could not be found")
|
||||
|
||||
def _file_item_exists(self, file_item_uuid: str) -> Union[FileSystemFile, FileSystemFolder, None]:
|
||||
"""Returns true if the file or folder UUID exists."""
|
||||
item = next((x for x in self.files if x.get_item_uuid() == file_item_uuid), None)
|
||||
if item:
|
||||
return item
|
||||
file = src.get_file(file_id=file_id)
|
||||
if file is None:
|
||||
raise Exception(f"file with UUID {file_id} could not be found")
|
||||
|
||||
next((x for x in self.folders if x.get_item_uuid() == file_item_uuid), None)
|
||||
# remove file from src
|
||||
src.remove_file(file_id)
|
||||
|
||||
if item:
|
||||
return item
|
||||
# add file to target
|
||||
target.add_file(file)
|
||||
|
||||
raise Exception(f"No file or folder found with id: {file_item_uuid}")
|
||||
def copy_file(self, src_folder_id: str, target_folder_id: str, file_id: str):
|
||||
"""Copies a file from one folder to another."""
|
||||
# check that both folders and the file exists
|
||||
src = self.get_folder_by_id(src_folder_id)
|
||||
target = self.get_folder_by_id(target_folder_id)
|
||||
|
||||
if src is None:
|
||||
raise Exception(f"src folder with UUID {src_folder_id} could not be found")
|
||||
|
||||
if target is None:
|
||||
raise Exception(f"src folder with UUID {target_folder_id} could not be found")
|
||||
|
||||
file = src.get_file(file_id=file_id)
|
||||
if file is None:
|
||||
raise Exception(f"file with UUID {file_id} could not be found")
|
||||
|
||||
# add file to target
|
||||
target.add_file(file)
|
||||
|
||||
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)
|
||||
if file is not None:
|
||||
return file
|
||||
|
||||
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)
|
||||
|
||||
@@ -1,30 +1,43 @@
|
||||
from typing import Dict
|
||||
|
||||
from pydantic import PrivateAttr
|
||||
|
||||
from primaite.simulator.core import SimComponent
|
||||
from primaite.simulator.file_system.file_system_file_type import FileSystemFileType
|
||||
from primaite.simulator.file_system.file_system_item_abc import FileSystemItemABC
|
||||
|
||||
|
||||
class FileSystemFile(FileSystemItemABC):
|
||||
class FileSystemFile(SimComponent):
|
||||
"""Class that represents a file in the simulation."""
|
||||
|
||||
_file_type: FileSystemFileType
|
||||
_file_type: FileSystemFileType = PrivateAttr()
|
||||
"""The type of the FileSystemFile"""
|
||||
|
||||
_file_size: float = PrivateAttr()
|
||||
"""Disk size of the FileSystemItem"""
|
||||
|
||||
def __init__(self, file_type: FileSystemFileType, file_size: float, **kwargs):
|
||||
"""
|
||||
Initialise FileSystemFile class.
|
||||
|
||||
:param item_parent: The UUID of the FileSystemItem parent
|
||||
:type item_parent: str
|
||||
|
||||
:param file_size: The size of the FileSystemItem
|
||||
:type file_size: float
|
||||
"""
|
||||
super().__init__(**kwargs)
|
||||
|
||||
self._file_type = file_type
|
||||
self._file_size = file_size
|
||||
|
||||
def get_file_size(self) -> float:
|
||||
"""Returns the size of the file system item."""
|
||||
return self._file_size
|
||||
|
||||
def get_file_type(self) -> FileSystemFileType:
|
||||
"""Returns the FileSystemFileType of the file."""
|
||||
return self._file_type
|
||||
|
||||
def move(self, target_directory: str):
|
||||
"""
|
||||
Changes the parent_item of the FileSystemFile.
|
||||
|
||||
Essentially simulates the file system item being moved from folder to folder
|
||||
|
||||
:param target_directory: The UUID of the directory the file system item should be moved to
|
||||
:type target_directory: str
|
||||
"""
|
||||
super().move(target_directory)
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Get the current state of the FileSystemFile as a dict.
|
||||
|
||||
@@ -1,14 +1,50 @@
|
||||
from typing import Dict
|
||||
from typing import Dict, List
|
||||
|
||||
from primaite.simulator.file_system.file_system_item_abc import FileSystemItemABC
|
||||
from pydantic import PrivateAttr
|
||||
|
||||
from primaite.simulator.core import SimComponent
|
||||
from primaite.simulator.file_system.file_system_file import FileSystemFile
|
||||
|
||||
|
||||
class FileSystemFolder(FileSystemItemABC):
|
||||
class FileSystemFolder(SimComponent):
|
||||
"""Simulation FileSystemFolder."""
|
||||
|
||||
_is_quarantined: bool
|
||||
_files: List[FileSystemFile] = PrivateAttr([])
|
||||
"""List of files stored in the folder."""
|
||||
|
||||
_folder_size: float = PrivateAttr(0)
|
||||
"""The current size of the folder"""
|
||||
|
||||
_is_quarantined: bool = PrivateAttr(False)
|
||||
"""Flag that marks the folder as quarantined if true."""
|
||||
|
||||
def get_files(self) -> List[FileSystemFile]:
|
||||
"""Returns the list of files the folder contains."""
|
||||
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)
|
||||
|
||||
def add_file(self, file: FileSystemFile):
|
||||
"""Adds a file to the folder list."""
|
||||
self._folder_size += file.get_file_size()
|
||||
|
||||
# add to list
|
||||
self._files.append(file)
|
||||
|
||||
def remove_file(self, file_id: str):
|
||||
"""Removes a file from the folder list."""
|
||||
file = next((f for f in self._files if f.uuid == file_id), None)
|
||||
self._files.remove(file)
|
||||
|
||||
# remove folder size from folder
|
||||
self._folder_size -= file.get_file_size()
|
||||
|
||||
def get_folder_size(self) -> float:
|
||||
"""Returns a sum of all file sizes in the files list."""
|
||||
return sum([file.get_file_size() for file in self._files])
|
||||
|
||||
def quarantine(self):
|
||||
"""Quarantines the File System Folder."""
|
||||
self._is_quarantined = True
|
||||
@@ -21,17 +57,6 @@ class FileSystemFolder(FileSystemItemABC):
|
||||
"""Returns true if the folder is being quarantined."""
|
||||
return self._is_quarantined
|
||||
|
||||
def move(self, target_directory: str):
|
||||
"""
|
||||
Changes the parent_item of the file system item.
|
||||
|
||||
Essentially simulates the file system item being moved from folder to folder
|
||||
|
||||
:param target_directory: The UUID of the directory the file system item should be moved to
|
||||
:type target_directory: str
|
||||
"""
|
||||
super().move(target_directory)
|
||||
|
||||
def describe_state(self) -> Dict:
|
||||
"""
|
||||
Get the current state of the FileSystemFolder as a dict.
|
||||
|
||||
@@ -1,60 +0,0 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from uuid import uuid4
|
||||
|
||||
from primaite.simulator.core import SimComponent
|
||||
|
||||
|
||||
class FileSystemItemABC(SimComponent, ABC):
|
||||
"""Abstract Base class for any file system items e.g. files and folders."""
|
||||
|
||||
_uuid: str
|
||||
"""Unique identifier for the FileSystemItem"""
|
||||
|
||||
_parent_item: str
|
||||
"""UUID of the parent FileSystemItem"""
|
||||
|
||||
_item_size: float
|
||||
"""Disk size of the FileSystemItem"""
|
||||
|
||||
def __init__(self, parent_item: str, item_size: float):
|
||||
"""
|
||||
Abstract base class used by FileSystem items.
|
||||
|
||||
:param parent_item: The UUID of the FileSystemItem parent
|
||||
:type parent_item: str
|
||||
|
||||
:param item_size: The size of the FileSystemItem
|
||||
:type item_size: float
|
||||
"""
|
||||
super().__init__()
|
||||
|
||||
# generate random uuid for file system item
|
||||
self._uuid = str(uuid4())
|
||||
|
||||
self._parent_item = parent_item
|
||||
|
||||
self._item_size = item_size
|
||||
|
||||
def get_item_uuid(self) -> str:
|
||||
"""Returns the file system item UUID."""
|
||||
return self._uuid
|
||||
|
||||
def get_item_parent(self) -> str:
|
||||
"""Returns the UUID of the item's parent."""
|
||||
return self._parent_item
|
||||
|
||||
def get_item_size(self) -> float:
|
||||
"""Returns the item size."""
|
||||
return self._item_size
|
||||
|
||||
@abstractmethod
|
||||
def move(self, target_directory: str):
|
||||
"""
|
||||
Changes the parent_item of the file system item.
|
||||
|
||||
Essentially simulates the file system item being moved from folder to folder
|
||||
|
||||
:param target_directory: The UUID of the directory the file system item should be moved to
|
||||
:type target_directory: str
|
||||
"""
|
||||
self._parent_item = target_directory
|
||||
@@ -0,0 +1,80 @@
|
||||
from primaite.simulator.file_system.file_system import FileSystem
|
||||
|
||||
|
||||
def test_create_folder_and_file():
|
||||
"""Test creating a folder and a file."""
|
||||
file_system = FileSystem()
|
||||
folder = file_system.create_folder()
|
||||
assert len(file_system.get_folders()) is 1
|
||||
|
||||
file_system.create_file(file_size=10, folder_uuid=folder.uuid)
|
||||
assert len(file_system.get_folders()[0].get_files()) is 1
|
||||
|
||||
|
||||
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_size=10)
|
||||
assert len(file_system.get_folders()) is 1
|
||||
assert file_system.get_folders()[0].get_file(file.uuid) is file
|
||||
|
||||
|
||||
def test_delete_file():
|
||||
"""Tests that a file can be deleted."""
|
||||
file_system = FileSystem()
|
||||
|
||||
file = file_system.create_file(file_size=10)
|
||||
assert len(file_system.get_folders()) is 1
|
||||
assert file_system.get_folders()[0].get_file(file.uuid) is file
|
||||
|
||||
file_system.delete_file(file.uuid)
|
||||
assert len(file_system.get_folders()) is 1
|
||||
assert len(file_system.get_folders()[0].get_files()) is 0
|
||||
|
||||
|
||||
def test_delete_folder():
|
||||
file_system = FileSystem()
|
||||
folder = file_system.create_folder()
|
||||
assert len(file_system.get_folders()) is 1
|
||||
|
||||
file_system.delete_folder(folder.uuid)
|
||||
assert len(file_system.get_folders()) is 0
|
||||
|
||||
|
||||
def test_move_file():
|
||||
"""Tests the file move function."""
|
||||
file_system = FileSystem()
|
||||
src_folder = file_system.create_folder()
|
||||
assert len(file_system.get_folders()) is 1
|
||||
|
||||
target_folder = file_system.create_folder()
|
||||
assert len(file_system.get_folders()) is 2
|
||||
|
||||
file = file_system.create_file(file_size=10, folder_uuid=src_folder.uuid)
|
||||
assert len(file_system.get_folder_by_id(src_folder.uuid).get_files()) is 1
|
||||
assert len(file_system.get_folder_by_id(target_folder.uuid).get_files()) is 0
|
||||
|
||||
file_system.move_file(src_folder.uuid, target_folder.uuid, file.uuid)
|
||||
|
||||
assert len(file_system.get_folder_by_id(src_folder.uuid).get_files()) is 0
|
||||
assert len(file_system.get_folder_by_id(target_folder.uuid).get_files()) is 1
|
||||
|
||||
|
||||
def test_copy_file():
|
||||
"""Tests the file copy function."""
|
||||
file_system = FileSystem()
|
||||
src_folder = file_system.create_folder()
|
||||
assert len(file_system.get_folders()) is 1
|
||||
|
||||
target_folder = file_system.create_folder()
|
||||
assert len(file_system.get_folders()) is 2
|
||||
|
||||
file = file_system.create_file(file_size=10, folder_uuid=src_folder.uuid)
|
||||
assert len(file_system.get_folder_by_id(src_folder.uuid).get_files()) is 1
|
||||
assert len(file_system.get_folder_by_id(target_folder.uuid).get_files()) is 0
|
||||
|
||||
file_system.copy_file(src_folder.uuid, target_folder.uuid, file.uuid)
|
||||
|
||||
assert len(file_system.get_folder_by_id(src_folder.uuid).get_files()) is 1
|
||||
assert len(file_system.get_folder_by_id(target_folder.uuid).get_files()) is 1
|
||||
|
||||
@@ -0,0 +1,14 @@
|
||||
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(file_size=1.5, file_type=FileSystemFileType.TBD)
|
||||
assert file.get_file_type() is FileSystemFileType.TBD
|
||||
|
||||
|
||||
def test_get_file_size():
|
||||
"""Tests that the file size is being returned properly."""
|
||||
file = FileSystemFile(file_size=1.5, file_type=FileSystemFileType.TBD)
|
||||
assert file.get_file_size() is 1.5
|
||||
|
||||
@@ -0,0 +1,41 @@
|
||||
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():
|
||||
folder = FileSystemFolder()
|
||||
|
||||
file = FileSystemFile(file_size=10, file_type=FileSystemFileType.TBD)
|
||||
|
||||
folder.add_file(file)
|
||||
assert folder.get_folder_size() is 10
|
||||
assert len(folder.get_files()) is 1
|
||||
|
||||
folder.remove_file(file_id=file.uuid)
|
||||
assert folder.get_folder_size() is 0
|
||||
assert len(folder.get_files()) is 0
|
||||
|
||||
|
||||
def test_get_file_by_id():
|
||||
folder = FileSystemFolder()
|
||||
|
||||
file = FileSystemFile(file_size=10, file_type=FileSystemFileType.TBD)
|
||||
|
||||
folder.add_file(file)
|
||||
assert folder.get_folder_size() is 10
|
||||
assert len(folder.get_files()) is 1
|
||||
|
||||
assert folder.get_file(file_id=file.uuid) is file
|
||||
|
||||
|
||||
def test_folder_quarantine_state():
|
||||
folder = FileSystemFolder()
|
||||
|
||||
assert folder.quarantine_status() is False
|
||||
|
||||
folder.quarantine()
|
||||
assert folder.quarantine_status() is True
|
||||
|
||||
folder.end_quarantine()
|
||||
assert folder.quarantine_status() is False
|
||||
|
||||
Reference in New Issue
Block a user