oxbrbtx/PrimAITE
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

#1714: update to use objects instead of uuids + tests

Browse Source
This commit is contained in:
Czar Echavez
2023-08-07 09:34:59 +01:00
parent 46c70ac084
commit 028211d288
8 changed files with 211 additions and 104 deletions
Download Patch File Download Diff File Expand all files Collapse all files

158
src/primaite/simulator/file_system/file_system.py
View File

@@ -1,3 +1,4 @@
from random import choice
from typing import Dict, List, Optional
from primaite.simulator.core import SimComponent
@@ -24,95 +25,148 @@ class FileSystem(SimComponent):
"""Returns the list of folders."""
return self.folders
def create_file(self, file_size: float, folder_uuid: Optional[str] = None) -> FileSystemFile:
def create_file(
self,
file_name: str,
file_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 file_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: file_size: The size the file takes on disk.
:type: file_size: Optional[float]
:param: file_type: The type of the file
:type: Optional[FileSystemFileType]
:param: folder_uuid: The uuid of the folder to add the file to
:type: folder_uuid: str
:type: folder_uuid: Optional[str]
"""
file = None
# if no folder uuid provided, create a folder and add file to it
if folder_uuid is None:
folder = FileSystemFolder()
folder = None
file = FileSystemFile(item_parent=folder.uuid, file_size=file_size, file_type=FileSystemFileType.TBD)
folder.add_file(file)
self.folders.append(folder)
else:
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(file_size=file_size, file_type=FileSystemFileType.TBD)
folder.add_file(file=file)
if folder is not None:
file = FileSystemFile(item_name=file_name, item_size=file_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(item_name="root")
# 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)
return file
def create_folder(self) -> FileSystemFolder:
def create_folder(
self,
folder_name: str,
) -> FileSystemFolder:
"""Creates a FileSystemFolder and adds it to the list of folders."""
folder = FileSystemFolder(item_parent=None)
folder = FileSystemFolder(item_name=folder_name)
self.folders.append(folder)
return folder
def delete_file(self, file_id: str):
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]
:param file_id: The UUID of the file item to delete
:type file_id: str
: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_id)
folder.remove_file(file=file)
def delete_folder(self, folder_id: str):
def delete_folder(self, folder: FileSystemFolder):
"""
Deletes a folder, removes it from the folders list and removes any child folders and files.
:param folder_id: The UUID of the file item to delete
:type folder_id: str
:param folder: The folder to remove
:type folder: FileSystemFolder
"""
self.folders = list(filter(lambda f: (f.uuid != folder_id), self.folders))
self.folders.remove(folder)
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)
def move_file(self, file: FileSystemFile, src_folder: FileSystemFolder, target_folder: FileSystemFolder):
"""
Moves a file from one folder to another.
if src is None:
raise Exception(f"src folder with UUID {src_folder_id} could not be found")
can provide
if target is None:
raise Exception(f"src folder with UUID {target_folder_id} could not be found")
: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")
file = src.get_file(file_id=file_id)
if file is None:
raise Exception(f"file with UUID {file_id} could not be found")
raise Exception("File to be moved is None")
# remove file from src
src.remove_file(file_id)
src_folder.remove_file(file)
# add file to target
target.add_file(file)
target_folder.add_file(file)
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)
def copy_file(self, file: FileSystemFile, src_folder: FileSystemFolder, target_folder: FileSystemFolder):
"""
Copies a file from one folder to another.
if src is None:
raise Exception(f"src folder with UUID {src_folder_id} could not be found")
can provide
if target is None:
raise Exception(f"src folder with UUID {target_folder_id} could not be found")
: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")
file = src.get_file(file_id=file_id)
if file is None:
raise Exception(f"file with UUID {file_id} could not be found")
raise Exception("File to be moved is None")
# add file to target
target.add_file(file)
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."""
@@ -121,6 +175,18 @@ class FileSystem(SimComponent):
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
"""
return next((f for f in self.folders if f.get_folder_name() == folder_name), None)
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)
def get_random_file_type(self) -> FileSystemFileType:
"""Returns a random FileSystemFileTypeEnum."""
return choice(list(FileSystemFileType))

37
src/primaite/simulator/file_system/file_system_file.py
View File

@@ -1,45 +1,54 @@
from random import choice, random
from random import choice, randint
from typing import Dict
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(SimComponent):
class FileSystemFile(FileSystemItemABC):
"""Class that represents a file in the simulation."""
file_type: FileSystemFileType = None
"""The type of the FileSystemFile"""
file_size: float = 0
"""Disk size of the FileSystemItem"""
def __init__(self, **kwargs):
"""
Initialise FileSystemFile class.
:param item_name: The name of the file.
:type item_name: str
:param file_type: The FileSystemFileType of the file
:type file_type: Optional[FileSystemFileType]
:param file_size: The size of the FileSystemItem
:type file_size: Optional[float]
:param item_size: The size of the FileSystemItem
:type item_size: Optional[float]
"""
super().__init__(**kwargs)
self.file_type = choice(list(FileSystemFileType))
self.file_size = random()
# set random file type if none provided
if kwargs.get("item_name") is None:
raise Exception("File name not provided.")
# set random file size if non provided
if kwargs.get("file_size") is not None:
self.file_size = kwargs.get("file_size")
if kwargs.get("item_size") is None:
kwargs["item_size"] = float(randint(1, 1000))
# set random file type if none provided
if kwargs.get("file_type") is None:
self.file_type = kwargs.get("file_type")
kwargs["file_type"] = choice(list(FileSystemFileType))
self.item_name = kwargs.get("item_name")
self.item_size = kwargs.get("item_size")
self.file_type = kwargs.get("file_type")
def get_file_name(self) -> str:
"""Returns the name of the file."""
return self.item_name
def get_file_size(self) -> float:
"""Returns the size of the file system item."""
return self.file_size
return self.item_size
def get_file_type(self) -> FileSystemFileType:
"""Returns the FileSystemFileType of the file."""

8
src/primaite/simulator/file_system/file_system_file_type.py
View File

@@ -4,4 +4,10 @@ from enum import Enum
class FileSystemFileType(str, Enum):
"""Enum used to determine the FileSystemFile type."""
TBD = "TBD"
CSV = ("CSV",)
DOC = ("DOC",)
EXE = ("EXE",)
PDF = ("PDF",)
TXT = ("TXT",)
XML = ("XML",)
ZIP = "ZIP"

33
src/primaite/simulator/file_system/file_system_folder.py
View File

@@ -1,21 +1,26 @@
from typing import Dict, List
from typing import Dict, List, Optional
from primaite.simulator.core import SimComponent
from primaite.simulator.file_system.file_system_file import FileSystemFile
from primaite.simulator.file_system.file_system_item_abc import FileSystemItemABC
class FileSystemFolder(SimComponent):
class FileSystemFolder(FileSystemItemABC):
"""Simulation FileSystemFolder."""
files: List[FileSystemFile] = []
"""List of files stored in the folder."""
folder_size: float = 0
"""The current size of the folder"""
is_quarantined: bool = False
"""Flag that marks the folder as quarantined if true."""
def get_folder_name(self) -> str:
"""Returns the item_name of the folder."""
return self.item_name
def get_folder_size(self) -> float:
"""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."""
return self.files
@@ -26,18 +31,24 @@ class FileSystemFolder(SimComponent):
def add_file(self, file: FileSystemFile):
"""Adds a file to the folder list."""
self.folder_size += file.get_file_size()
self.item_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)
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]
"""
self.files.remove(file)
# remove folder size from folder
self.folder_size -= file.get_file_size()
self.item_size -= file.get_file_size()
def get_folder_size(self) -> float:
"""Returns a sum of all file sizes in the files list."""

13
src/primaite/simulator/file_system/file_system_item_abc.py Normal file
View File

@@ -0,0 +1,13 @@
from abc import ABC
from primaite.simulator.core import SimComponent
class FileSystemItemABC(SimComponent, ABC):
"""Abstract base class for FileSystemItems used in the file system simulation."""
item_size: float = 0
"""The size the item takes up on disk."""
item_name: str
"""The name of the file system item."""

36
tests/unit_tests/_primaite/_simulator/_file_system/test_file_system.py
View File

@@ -4,18 +4,20 @@ 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()
folder = file_system.create_folder(folder_name="test_folder")
assert len(file_system.get_folders()) is 1
file_system.create_file(file_size=10, folder_uuid=folder.uuid)
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() is 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_size=10)
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
@@ -24,38 +26,38 @@ def test_delete_file():
"""Tests that a file can be deleted."""
file_system = FileSystem()
file = file_system.create_file(file_size=10)
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
file_system.delete_file(file.uuid)
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
def test_delete_folder():
file_system = FileSystem()
folder = file_system.create_folder()
folder = file_system.create_folder(folder_name="test_folder")
assert len(file_system.get_folders()) is 1
file_system.delete_folder(folder.uuid)
file_system.delete_folder(folder)
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()
src_folder = file_system.create_folder(folder_name="test_folder_1")
assert len(file_system.get_folders()) is 1
target_folder = file_system.create_folder()
target_folder = file_system.create_folder(folder_name="test_folder_2")
assert len(file_system.get_folders()) is 2
file = file_system.create_file(file_size=10, folder_uuid=src_folder.uuid)
file = file_system.create_file(file_name="test_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)
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).get_files()) is 0
assert len(file_system.get_folder_by_id(target_folder.uuid).get_files()) is 1
@@ -64,17 +66,17 @@ def test_move_file():
def test_copy_file():
"""Tests the file copy function."""
file_system = FileSystem()
src_folder = file_system.create_folder()
src_folder = file_system.create_folder(folder_name="test_folder_1")
assert len(file_system.get_folders()) is 1
target_folder = file_system.create_folder()
target_folder = file_system.create_folder(folder_name="test_folder_2")
assert len(file_system.get_folders()) is 2
file = file_system.create_file(file_size=10, folder_uuid=src_folder.uuid)
file = file_system.create_file(file_name="test_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)
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).get_files()) is 1
assert len(file_system.get_folder_by_id(target_folder.uuid).get_files()) is 1
@@ -83,10 +85,10 @@ def test_copy_file():
def test_serialisation():
"""Test to check that the object serialisation works correctly."""
file_system = FileSystem()
folder = file_system.create_folder()
folder = file_system.create_folder(folder_name="test_folder")
assert len(file_system.get_folders()) is 1
file_system.create_file(file_size=10, folder_uuid=folder.uuid)
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
serialised_file_sys = file_system.model_dump_json()

10
tests/unit_tests/_primaite/_simulator/_file_system/test_file_system_file.py
View File

@@ -4,20 +4,20 @@ from primaite.simulator.file_system.file_system_file_type import FileSystemFileT
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
file = FileSystemFile(item_name="test", file_type=FileSystemFileType.DOC)
assert file.get_file_type() is FileSystemFileType.DOC
def test_get_file_size():
"""Tests that the file size is being returned properly."""
file = FileSystemFile(file_size=1.5, file_type=FileSystemFileType.TBD)
file = FileSystemFile(item_name="test", item_size=1.5)
assert file.get_file_size() is 1.5
def test_serialisation():
"""Test to check that the object serialisation works correctly."""
file = FileSystemFile(file_size=1.5, file_type=FileSystemFileType.TBD)
file = FileSystemFile(item_name="test", item_size=1.5, file_type=FileSystemFileType.DOC)
serialised_file = file.model_dump_json()
deserialised_file = FileSystemFile().model_validate_json(serialised_file)
deserialised_file = FileSystemFile(item_name="test").model_validate_json(serialised_file)
assert file == deserialised_file

20
tests/unit_tests/_primaite/_simulator/_file_system/test_file_system_folder.py
View File

@@ -5,25 +5,25 @@ 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()
folder = FileSystemFolder(item_name="test")
file = FileSystemFile(file_size=10, file_type=FileSystemFileType.TBD)
file = FileSystemFile(item_name="test_file", item_size=10, file_type=FileSystemFileType.DOC)
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)
folder.remove_file(file)
assert folder.get_folder_size() is 0
assert len(folder.get_files()) is 0
def test_get_file_by_id():
"""Test to make sure that the correct file is returned."""
folder = FileSystemFolder()
folder = FileSystemFolder(item_name="test")
file = FileSystemFile(file_size=10, file_type=FileSystemFileType.TBD)
file2 = FileSystemFile(file_size=10, file_type=FileSystemFileType.TBD)
file = FileSystemFile(item_name="test_file", item_size=10, file_type=FileSystemFileType.DOC)
file2 = FileSystemFile(item_name="test_file_2", item_size=10, file_type=FileSystemFileType.DOC)
folder.add_file(file)
folder.add_file(file2)
@@ -35,7 +35,7 @@ def test_get_file_by_id():
def test_folder_quarantine_state():
"""Tests the changing of folder quarantine status."""
folder = FileSystemFolder()
folder = FileSystemFolder(item_name="test")
assert folder.quarantine_status() is False
@@ -48,12 +48,12 @@ def test_folder_quarantine_state():
def test_serialisation():
"""Test to check that the object serialisation works correctly."""
folder = FileSystemFolder()
file = FileSystemFile(file_size=10, file_type=FileSystemFileType.TBD)
folder = FileSystemFolder(item_name="test")
file = FileSystemFile(item_name="test_file", item_size=10, file_type=FileSystemFileType.DOC)
folder.add_file(file)
serialised_folder = folder.model_dump_json()
deserialised_folder = FileSystemFolder().model_validate_json(serialised_folder)
deserialised_folder = FileSystemFolder(item_name="test").model_validate_json(serialised_folder)
assert folder == deserialised_folder