diff --git a/moonshot/src/api/api_bookmark.py b/moonshot/src/api/api_bookmark.py index bc4ecb87..bd09cebf 100644 --- a/moonshot/src/api/api_bookmark.py +++ b/moonshot/src/api/api_bookmark.py @@ -39,7 +39,7 @@ def api_insert_bookmark( metric=metric, bookmark_time="", # bookmark_time will be set to current time in add_bookmark method ) - return Bookmark.get_instance().add_bookmark(bookmark_args) + return Bookmark().add_bookmark(bookmark_args) def api_get_all_bookmarks() -> list[dict]: @@ -49,7 +49,7 @@ def api_get_all_bookmarks() -> list[dict]: Returns: list[dict]: A list of bookmarks, each represented as a dictionary. """ - return Bookmark.get_instance().get_all_bookmarks() + return Bookmark().get_all_bookmarks() def api_get_bookmark(bookmark_name: str) -> dict: @@ -62,7 +62,7 @@ def api_get_bookmark(bookmark_name: str) -> dict: Returns: dict: The bookmark details corresponding to the provided ID. """ - return Bookmark.get_instance().get_bookmark(bookmark_name) + return Bookmark().get_bookmark(bookmark_name) def api_delete_bookmark(bookmark_name: str) -> dict: @@ -72,14 +72,14 @@ def api_delete_bookmark(bookmark_name: str) -> dict: Args: bookmark_name (str): The name of the bookmark to be removed. """ - return Bookmark.get_instance().delete_bookmark(bookmark_name) + return Bookmark().delete_bookmark(bookmark_name) def api_delete_all_bookmark() -> dict: """ Removes all bookmarks from the database. """ - return Bookmark.get_instance().delete_all_bookmark() + return Bookmark().delete_all_bookmark() def api_export_bookmarks(export_file_name: str = "bookmarks") -> str: @@ -92,4 +92,4 @@ def api_export_bookmarks(export_file_name: str = "bookmarks") -> str: Returns: str: The filepath of where the file is written. """ - return Bookmark.get_instance().export_bookmarks(export_file_name) + return Bookmark().export_bookmarks(export_file_name) diff --git a/moonshot/src/bookmark/bookmark.py b/moonshot/src/bookmark/bookmark.py index dd8a1556..42376522 100644 --- a/moonshot/src/bookmark/bookmark.py +++ b/moonshot/src/bookmark/bookmark.py @@ -1,7 +1,24 @@ +from __future__ import annotations + +import textwrap from datetime import datetime from moonshot.src.bookmark.bookmark_arguments import BookmarkArguments from moonshot.src.configs.env_variables import EnvVariables +from moonshot.src.messages_constants import ( + BOOKMARK_ADD_BOOKMARK_ERROR, + BOOKMARK_ADD_BOOKMARK_SUCCESS, + BOOKMARK_ADD_BOOKMARK_VALIDATION_ERROR, + BOOKMARK_DELETE_ALL_BOOKMARK_ERROR, + BOOKMARK_DELETE_ALL_BOOKMARK_SUCCESS, + BOOKMARK_DELETE_BOOKMARK_ERROR, + BOOKMARK_DELETE_BOOKMARK_ERROR_1, + BOOKMARK_DELETE_BOOKMARK_SUCCESS, + BOOKMARK_EXPORT_BOOKMARK_ERROR, + BOOKMARK_EXPORT_BOOKMARK_VALIDATION_ERROR, + BOOKMARK_GET_BOOKMARK_ERROR, + BOOKMARK_GET_BOOKMARK_ERROR_1, +) from moonshot.src.storage.storage import Storage from moonshot.src.utils.log import configure_logger @@ -12,36 +29,7 @@ class Bookmark: _instance = None - def __new__(cls, db_name="bookmark"): - """ - Create a new instance of the Bookmark class or return the existing instance. - - Args: - db_name (str): The name of the database. - - Returns: - Bookmark: The singleton instance of the Bookmark class. - """ - if cls._instance is None: - cls._instance = super(Bookmark, cls).__new__(cls) - cls._instance.__init_instance(db_name) - return cls._instance - - @classmethod - def get_instance(cls, db_name="bookmark"): - """ - Get the singleton instance of the Bookmark class. - - Args: - db_name (str): The name of the database. - - Returns: - Bookmark: The singleton instance of the Bookmark class. - """ - if cls._instance is None: - cls._instance = super(Bookmark, cls).__new__(cls) - cls._instance.__init_instance(db_name) - return cls._instance + sql_table_name = "bookmark" sql_create_bookmark_table = """ CREATE TABLE IF NOT EXISTS bookmark ( @@ -77,19 +65,46 @@ def get_instance(cls, db_name="bookmark"): DELETE FROM bookmark; """ - def __init_instance(self, db_name) -> None: + def __new__(cls, db_name="bookmark"): + """ + Create a new instance of the Bookmark class or return the existing instance. + + This method ensures that only one instance of the Bookmark class is created (singleton pattern). + If an instance already exists, it returns that instance. Otherwise, it creates a new instance + and initializes it with the provided database name. + + Args: + db_name (str): The name of the database. Defaults to "bookmark". + + Returns: + Bookmark: The singleton instance of the Bookmark class. + """ + if cls._instance is None: + cls._instance = super(Bookmark, cls).__new__(cls) + cls._instance.__init_instance(db_name) + return cls._instance + + def __init_instance(self, db_name: str = "bookmark") -> None: """ Initialize the database instance for the Bookmark class. + This method sets up the database connection for the Bookmark class. It creates a new database + connection using the provided database name and checks if the required table exists. If the table + does not exist, it creates the table. + Args: - db_name (str): The name of the database. + db_name (str): The name of the database. Defaults to "bookmark". """ self.db_instance = Storage.create_database_connection( EnvVariables.BOOKMARKS.name, db_name, "db" ) - Storage.create_database_table( - self.db_instance, Bookmark.sql_create_bookmark_table - ) + + if not Storage.check_database_table_exists( + self.db_instance, Bookmark.sql_table_name + ): + Storage.create_database_table( + self.db_instance, Bookmark.sql_create_bookmark_table + ) def add_bookmark(self, bookmark: BookmarkArguments) -> dict: """ @@ -99,7 +114,7 @@ def add_bookmark(self, bookmark: BookmarkArguments) -> dict: bookmark (BookmarkArguments): The bookmark data to add. Returns: - bool: True if the bookmark was added successfully, False otherwise. + dict: A dictionary containing the success status and a message. """ bookmark.bookmark_time = datetime.now().replace(microsecond=0).isoformat(" ") @@ -119,12 +134,14 @@ def add_bookmark(self, bookmark: BookmarkArguments) -> dict: self.db_instance, data, Bookmark.sql_insert_bookmark_record ) if results is not None: - return {"success": True, "message": "Bookmark added successfully."} + return {"success": True, "message": BOOKMARK_ADD_BOOKMARK_SUCCESS} else: - raise Exception("Error inserting record into database.") + raise Exception(BOOKMARK_ADD_BOOKMARK_VALIDATION_ERROR) except Exception as e: - error_message = f"Failed to add bookmark record: {e}" - return {"success": False, "message": error_message} + return { + "success": False, + "message": BOOKMARK_ADD_BOOKMARK_ERROR.format(message=str(e)), + } def get_all_bookmarks(self) -> list[dict]: """ @@ -137,7 +154,9 @@ def get_all_bookmarks(self) -> list[dict]: self.db_instance, Bookmark.sql_select_bookmarks_record, ) - if list_of_bookmarks_tuples: + if isinstance(list_of_bookmarks_tuples, list) and all( + isinstance(item, tuple) for item in list_of_bookmarks_tuples + ): list_of_bookmarks = [ BookmarkArguments.from_tuple_to_dict(bookmark_tuple) for bookmark_tuple in list_of_bookmarks_tuples @@ -151,7 +170,7 @@ def get_bookmark(self, bookmark_name: str) -> dict: Retrieve a bookmark by its unique name. Args: - bookmark_name (int): The unique name for the bookmark. + bookmark_name (str): The unique name for the bookmark. Returns: dict: The bookmark information as a dictionary. @@ -159,18 +178,24 @@ def get_bookmark(self, bookmark_name: str) -> dict: Raises: RuntimeError: If the bookmark cannot be found. """ - if bookmark_name is not None: + if isinstance(bookmark_name, str) and bookmark_name: bookmark_info = Storage.read_database_record( self.db_instance, (bookmark_name,), Bookmark.sql_select_bookmark_record ) - if bookmark_info is not None: + if ( + bookmark_info is not None + and isinstance(bookmark_info, tuple) + and all(isinstance(item, str) for item in bookmark_info) + ): return BookmarkArguments.from_tuple_to_dict(bookmark_info) else: raise RuntimeError( - f"[Bookmark] No record found for bookmark name {bookmark_name}" + BOOKMARK_GET_BOOKMARK_ERROR.format(message=bookmark_name) ) else: - raise RuntimeError(f"[Bookmark] Invalid bookmark name: {bookmark_name}") + raise RuntimeError( + BOOKMARK_GET_BOOKMARK_ERROR_1.format(message=bookmark_name) + ) def delete_bookmark(self, bookmark_name: str) -> dict: """ @@ -178,22 +203,33 @@ def delete_bookmark(self, bookmark_name: str) -> dict: Args: bookmark_name (str): The unique name for the bookmark to be deleted. + + Returns: + dict: A dictionary containing the success status and a message. """ - if bookmark_name is not None: + if isinstance(bookmark_name, str) and bookmark_name: try: - sql_delete_bookmark_record = f""" + sql_delete_bookmark_record = textwrap.dedent( + f""" DELETE FROM bookmark WHERE name = '{bookmark_name}'; """ + ) Storage.delete_database_record_in_table( self.db_instance, sql_delete_bookmark_record ) - return {"success": True, "message": "Bookmark record deleted."} + return {"success": True, "message": BOOKMARK_DELETE_BOOKMARK_SUCCESS} except Exception as e: - error_message = f"Failed to delete bookmark record: {e}" - return {"success": False, "message": error_message} + return { + "success": False, + "message": BOOKMARK_DELETE_BOOKMARK_ERROR.format(message=str(e)), + } else: - error_message = f"[Bookmark] Invalid bookmark name: {bookmark_name}" - return {"success": False, "message": error_message} + return { + "success": False, + "message": BOOKMARK_DELETE_BOOKMARK_ERROR_1.format( + message=bookmark_name + ), + } def delete_all_bookmark(self) -> dict: """ @@ -206,10 +242,12 @@ def delete_all_bookmark(self) -> dict: Storage.delete_database_record_in_table( self.db_instance, Bookmark.sql_delete_bookmark_records ) - return {"success": True, "message": "All bookmark records deleted."} + return {"success": True, "message": BOOKMARK_DELETE_ALL_BOOKMARK_SUCCESS} except Exception as e: - error_message = f"Failed to delete all bookmark records: {e}" - return {"success": False, "message": error_message} + return { + "success": False, + "message": BOOKMARK_DELETE_ALL_BOOKMARK_ERROR.format(message=str(e)), + } def export_bookmarks(self, export_file_name: str = "bookmarks") -> str: """ @@ -224,13 +262,27 @@ def export_bookmarks(self, export_file_name: str = "bookmarks") -> str: Returns: str: The path to the exported JSON file containing the bookmarks. + + Raises: + Exception: If the export file name is invalid or an error occurs during export. """ + if not isinstance(export_file_name, str) or not export_file_name: + error_message = BOOKMARK_EXPORT_BOOKMARK_ERROR.format( + message=BOOKMARK_EXPORT_BOOKMARK_VALIDATION_ERROR + ) + logger.error(error_message) + raise Exception(error_message) + list_of_bookmarks_tuples = Storage.read_database_records( self.db_instance, Bookmark.sql_select_bookmarks_record, ) - if list_of_bookmarks_tuples is not None: + if ( + list_of_bookmarks_tuples is not None + and isinstance(list_of_bookmarks_tuples, list) + and all(isinstance(item, tuple) for item in list_of_bookmarks_tuples) + ): bookmarks_json = [ BookmarkArguments.from_tuple_to_dict(bookmark_tuple) for bookmark_tuple in list_of_bookmarks_tuples @@ -246,12 +298,19 @@ def export_bookmarks(self, export_file_name: str = "bookmarks") -> str: "json", ) except Exception as e: - logger.error(f"Failed to export bookmarks - {str(e)}.") - raise e + error_message = BOOKMARK_EXPORT_BOOKMARK_ERROR.format(message=str(e)) + logger.error(error_message) + raise Exception(error_message) def close(self) -> None: """ - Close the database connection. + Close the database connection and set the Bookmark instance to None. + + This method ensures that the database connection is properly closed and the singleton + instance of the Bookmark class is reset to None, allowing for a fresh instance to be created + if needed in the future. """ if self.db_instance: Storage.close_database_connection(self.db_instance) + + Bookmark._instance = None diff --git a/moonshot/src/bookmark/bookmark_arguments.py b/moonshot/src/bookmark/bookmark_arguments.py index 3107dfce..4f3334c9 100644 --- a/moonshot/src/bookmark/bookmark_arguments.py +++ b/moonshot/src/bookmark/bookmark_arguments.py @@ -2,6 +2,10 @@ from pydantic import BaseModel, Field +from moonshot.src.messages_constants import ( + BOOKMARK_ARGUMENTS_FROM_TUPLE_TO_DICT_VALIDATION_ERROR, +) + class BookmarkArguments(BaseModel): name: str = Field(min_length=1) @@ -24,7 +28,13 @@ def from_tuple_to_dict(cls, values: tuple) -> dict: Returns: dict: A dictionary representing the BookmarkArguments. + + Raises: + ValueError: If the number of values in the tuple is less than 10. """ + if len(values) < 10: + raise ValueError(BOOKMARK_ARGUMENTS_FROM_TUPLE_TO_DICT_VALIDATION_ERROR) + return { "name": values[1], "prompt": values[2], diff --git a/moonshot/src/messages_constants.py b/moonshot/src/messages_constants.py new file mode 100644 index 00000000..cc76d6d1 --- /dev/null +++ b/moonshot/src/messages_constants.py @@ -0,0 +1,40 @@ +# ------------------------------------------------------------------------------ +# BOOKMARK - add_bookmark +# ------------------------------------------------------------------------------ +BOOKMARK_ADD_BOOKMARK_SUCCESS = "[Bookmark] Bookmark added successfully." +BOOKMARK_ADD_BOOKMARK_ERROR = "[Bookmark] Failed to add bookmark record: {message}" +BOOKMARK_ADD_BOOKMARK_VALIDATION_ERROR = "Error inserting record into database." + +# ------------------------------------------------------------------------------ +# BOOKMARK - get_bookmark +# ------------------------------------------------------------------------------ +BOOKMARK_GET_BOOKMARK_ERROR = "[Bookmark] No record found for bookmark name: {message}" +BOOKMARK_GET_BOOKMARK_ERROR_1 = "[Bookmark] Invalid bookmark name: {message}" + +# ------------------------------------------------------------------------------ +# BOOKMARK - delete_bookmark +# ------------------------------------------------------------------------------ +BOOKMARK_DELETE_BOOKMARK_SUCCESS = "[Bookmark] Bookmark record deleted." +BOOKMARK_DELETE_BOOKMARK_ERROR = ( + "[Bookmark] Failed to delete bookmark record: {message}" +) +BOOKMARK_DELETE_BOOKMARK_ERROR_1 = "[Bookmark] Invalid bookmark name: {message}" + +# ------------------------------------------------------------------------------ +# BOOKMARK - delete_all_bookmark +# ------------------------------------------------------------------------------ +BOOKMARK_DELETE_ALL_BOOKMARK_SUCCESS = "[Bookmark] All bookmark records deleted." +BOOKMARK_DELETE_ALL_BOOKMARK_ERROR = ( + "[Bookmark] Failed to delete all bookmark records: {message}" +) + +# ------------------------------------------------------------------------------ +# BOOKMARK - export_bookmarks +# ------------------------------------------------------------------------------ +BOOKMARK_EXPORT_BOOKMARK_ERROR = "[Bookmark] Failed to export bookmarks: {message}" +BOOKMARK_EXPORT_BOOKMARK_VALIDATION_ERROR = "Export filename must be a non-empty string" + +# ------------------------------------------------------------------------------ +# BOOKMARK ARGUMENTS - from_tuple_to_dict +# ------------------------------------------------------------------------------ +BOOKMARK_ARGUMENTS_FROM_TUPLE_TO_DICT_VALIDATION_ERROR = "[BookmarkArguments] Failed to convert to dictionary because of the insufficient number of values" # noqa: E501 diff --git a/tests/unit-tests/src/test_api_bookmark.py b/tests/unit-tests/src/test_api_bookmark.py index 73283e13..b933edd3 100644 --- a/tests/unit-tests/src/test_api_bookmark.py +++ b/tests/unit-tests/src/test_api_bookmark.py @@ -1,22 +1,21 @@ -import pytest import shutil -import os +from pathlib import Path +import pytest from pydantic import ValidationError -from unittest.mock import patch, MagicMock from moonshot.api import ( - api_get_bookmark, - api_export_bookmarks, + api_delete_all_bookmark, api_delete_bookmark, - api_insert_bookmark, + api_export_bookmarks, api_get_all_bookmarks, - api_delete_all_bookmark, + api_get_bookmark, + api_insert_bookmark, api_set_environment_variables, ) - from moonshot.src.bookmark.bookmark import Bookmark + class TestCollectionApiBookmark: @pytest.fixture(autouse=True) def init(self): @@ -29,32 +28,29 @@ def init(self): "IO_MODULES": "tests/unit-tests/src/data/io-modules/", } ) - - shutil.copyfile( + + # Copy the required bookmark database + shutil.copy( "tests/unit-tests/common/samples/bookmark.db", "tests/unit-tests/src/data/bookmarks/bookmark.db", ) + # Create a bookmark instance + bookmark_instance = Bookmark() + yield + # Terminate the bookmark_instance + bookmark_instance.close() + + # Delete bookmark database and accompanying files run_data_files = [ - # "tests/unit-tests/src/data/bookmarks/bookmark.db", - "tests/unit-tests/src/data/bookmarks/bookmark.json", + Path("tests/unit-tests/src/data/bookmarks/bookmark.db"), + Path("tests/unit-tests/src/data/bookmarks/bookmark.json"), ] for run_data_file in run_data_files: - if os.path.exists(run_data_file): - os.remove(run_data_file) - - # ------------------------------------------------------------------------------ - # Test bookmark instance is singleton - # ------------------------------------------------------------------------------ - def test_bookmark_singleton(self): - # Retrieve two instances of the Bookmark class - bookmark_instance_1 = Bookmark() - bookmark_instance_2 = Bookmark() - - # Assert that both instances are the same (singleton behavior) - assert bookmark_instance_1 is bookmark_instance_2, "Bookmark instances should be the same (singleton pattern)." + if run_data_file.exists(): + run_data_file.unlink() # ------------------------------------------------------------------------------ # Test api_insert_bookmark functionality @@ -62,29 +58,18 @@ def test_bookmark_singleton(self): @pytest.mark.parametrize( "input_args, expected_dict", [ - # Valid cases for bookmark insertions + # Valid cases for bookmark insertions # All Value present ( { "name": "Bookmark A", "prompt": "Prompt A", - "response": "Response A", "prepared_prompt": "Prepared Prompt", + "response": "Response A", "context_strategy": "Strategy A", "prompt_template": "Template A", - "attack_module": "Module A" - }, - {"expected_output": True}, - ), - # No Attack Module - ( - { - "name": "Bookmark B", - "prompt": "Prompt B", - "response": "Response B", - "prepared_prompt": "Prepared Prompt", - "context_strategy": "Strategy B", - "prompt_template": "Template B", + "attack_module": "Module A", + "metric": "Metric A", }, {"expected_output": True}, ), @@ -93,10 +78,11 @@ def test_bookmark_singleton(self): { "name": "Bookmark C", "prompt": "Prompt C", - "response": "Response C", "prepared_prompt": "Prepared Prompt", + "response": "Response C", "prompt_template": "Template C", - "attack_module": "Module C" + "attack_module": "Module C", + "metric": "Metric C", }, {"expected_output": True}, ), @@ -105,10 +91,36 @@ def test_bookmark_singleton(self): { "name": "Bookmark D", "prompt": "Prompt D", - "response": "Response D", "prepared_prompt": "Prepared Prompt", + "response": "Response D", "context_strategy": "Strategy D", - "attack_module": "Module D" + "attack_module": "Module D", + "metric": "Metric D", + }, + {"expected_output": True}, + ), + # No Attack Module + ( + { + "name": "Bookmark B", + "prompt": "Prompt B", + "prepared_prompt": "Prepared Prompt", + "response": "Response B", + "context_strategy": "Strategy B", + "prompt_template": "Template B", + "metric": "Metric B", + }, + {"expected_output": True}, + ), + # No Metric + ( + { + "name": "Bookmark A1", + "prompt": "Prompt A1", + "prepared_prompt": "Prepared Prompt", + "response": "Response A1", + "context_strategy": "Strategy A1", + "prompt_template": "Template A1", }, {"expected_output": True}, ), @@ -117,8 +129,8 @@ def test_bookmark_singleton(self): { "name": "Bookmark E", "prompt": "Prompt E", - "response": "Response E", "prepared_prompt": "Prepared Prompt", + "response": "Response E", }, {"expected_output": True}, ), @@ -127,8 +139,12 @@ def test_bookmark_singleton(self): { "name": None, "prompt": "Prompt F", - "response": "Response F", "prepared_prompt": "Prepared Prompt", + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -139,9 +155,13 @@ def test_bookmark_singleton(self): ( { "name": "", - "prompt": "Prompt G", - "response": "Response G", + "prompt": "Prompt F", "prepared_prompt": "Prepared Prompt", + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -149,12 +169,33 @@ def test_bookmark_singleton(self): "expected_exception": ValidationError, }, ), + ( + { + "name": (), + "prompt": "Prompt F", + "prepared_prompt": "Prepared Prompt", + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", + }, + { + "expected_output": False, + "expected_error_message": "Input should be a valid string", + "expected_exception": ValidationError, + }, + ), ( { "name": [], - "prompt": "Prompt H", - "response": "Response H", + "prompt": "Prompt F", "prepared_prompt": "Prepared Prompt", + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -165,9 +206,13 @@ def test_bookmark_singleton(self): ( { "name": {}, - "prompt": "Prompt I", - "response": "Response I", + "prompt": "Prompt F", "prepared_prompt": "Prepared Prompt", + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -178,9 +223,13 @@ def test_bookmark_singleton(self): ( { "name": 123, - "prompt": "Prompt J", - "response": "Response J", + "prompt": "Prompt F", "prepared_prompt": "Prepared Prompt", + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -193,8 +242,12 @@ def test_bookmark_singleton(self): { "name": "Bookmark K", "prompt": None, - "response": "Response K", "prepared_prompt": "Prepared Prompt", + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -206,8 +259,12 @@ def test_bookmark_singleton(self): { "name": "Bookmark L", "prompt": "", - "response": "Response L", "prepared_prompt": "Prepared Prompt", + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -215,12 +272,33 @@ def test_bookmark_singleton(self): "expected_exception": ValidationError, }, ), + ( + { + "name": "Bookmark L", + "prompt": (), + "prepared_prompt": "Prepared Prompt", + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", + }, + { + "expected_output": False, + "expected_error_message": "Input should be a valid string", + "expected_exception": ValidationError, + }, + ), ( { "name": "Bookmark M", "prompt": [], - "response": "Response M", "prepared_prompt": "Prepared Prompt", + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -232,8 +310,12 @@ def test_bookmark_singleton(self): { "name": "Bookmark N", "prompt": {}, - "response": "Response N", "prepared_prompt": "Prepared Prompt", + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -245,8 +327,12 @@ def test_bookmark_singleton(self): { "name": "Bookmark O", "prompt": 123, - "response": "Response O", "prepared_prompt": "Prepared Prompt", + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -254,13 +340,120 @@ def test_bookmark_singleton(self): "expected_exception": ValidationError, }, ), - # Invalid cases for 'prompt' + # Invalid cases for 'prepared_prompt' ( { "name": "Bookmark P", - "prompt": None, - "response": "Response P", + "prompt": "Prompt A", + "prepared_prompt": None, + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", + }, + { + "expected_output": False, + "expected_error_message": "Input should be a valid string", + "expected_exception": ValidationError, + }, + ), + ( + { + "name": "Bookmark P", + "prompt": "Prompt A", + "prepared_prompt": "", + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", + }, + { + "expected_output": False, + "expected_error_message": "String should have at least 1 character", + "expected_exception": ValidationError, + }, + ), + ( + { + "name": "Bookmark P", + "prompt": "Prompt A", + "prepared_prompt": (), + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", + }, + { + "expected_output": False, + "expected_error_message": "Input should be a valid string", + "expected_exception": ValidationError, + }, + ), + ( + { + "name": "Bookmark P", + "prompt": "Prompt A", + "prepared_prompt": [], + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", + }, + { + "expected_output": False, + "expected_error_message": "Input should be a valid string", + "expected_exception": ValidationError, + }, + ), + ( + { + "name": "Bookmark P", + "prompt": "Prompt A", + "prepared_prompt": {}, + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", + }, + { + "expected_output": False, + "expected_error_message": "Input should be a valid string", + "expected_exception": ValidationError, + }, + ), + ( + { + "name": "Bookmark P", + "prompt": "Prompt A", + "prepared_prompt": 123, + "response": "Response F", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", + }, + { + "expected_output": False, + "expected_error_message": "Input should be a valid string", + "expected_exception": ValidationError, + }, + ), + # Invalid cases for 'response' + ( + { + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", + "response": None, + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -270,10 +463,14 @@ def test_bookmark_singleton(self): ), ( { - "name": "Bookmark Q", - "prompt": "", - "response": "Response Q", + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", + "response": "", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -283,10 +480,14 @@ def test_bookmark_singleton(self): ), ( { - "name": "Bookmark R", - "prompt": [], - "response": "Response R", + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", + "response": (), + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -296,10 +497,14 @@ def test_bookmark_singleton(self): ), ( { - "name": "Bookmark S", - "prompt": {}, - "response": "Response S", + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", + "response": [], + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -309,10 +514,14 @@ def test_bookmark_singleton(self): ), ( { - "name": "Bookmark T", - "prompt": 123, - "response": "Response T", + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", + "response": {}, + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -320,14 +529,16 @@ def test_bookmark_singleton(self): "expected_exception": ValidationError, }, ), - - # Invalid cases for 'response' ( { - "name": "Bookmark U", - "prompt": "Prompt U", - "response": None, + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", + "response": 123, + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -335,25 +546,34 @@ def test_bookmark_singleton(self): "expected_exception": ValidationError, }, ), + # Invalid cases for 'context_strategy' ( { - "name": "Bookmark V", - "prompt": "Prompt V", - "response": "", + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", + "response": "Response A", + "context_strategy": None, + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, - "expected_error_message": "String should have at least 1 character", + "expected_error_message": "Input should be a valid string", "expected_exception": ValidationError, }, ), ( { - "name": "Bookmark W", - "prompt": "Prompt W", - "response": [], + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", + "response": "Response A", + "context_strategy": (), + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -363,10 +583,14 @@ def test_bookmark_singleton(self): ), ( { - "name": "Bookmark X", - "prompt": "Prompt X", - "response": {}, + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", + "response": "Response A", + "context_strategy": [], + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -376,10 +600,117 @@ def test_bookmark_singleton(self): ), ( { - "name": "Bookmark Y", - "prompt": "Prompt Y", - "response": 123, + "name": "Bookmark P", + "prompt": "Prompt A", + "prepared_prompt": "Prepared Prompt", + "response": "Response A", + "context_strategy": {}, + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", + }, + { + "expected_output": False, + "expected_error_message": "Input should be a valid string", + "expected_exception": ValidationError, + }, + ), + ( + { + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", + "response": "Response A", + "context_strategy": 123, + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", + }, + { + "expected_output": False, + "expected_error_message": "Input should be a valid string", + "expected_exception": ValidationError, + }, + ), + # Invalid cases for 'prompt_template' + ( + { + "name": "Bookmark P", + "prompt": "Prompt A", + "prepared_prompt": "Prepared Prompt", + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": None, + "attack_module": "Module A", + "metric": "Metric A", + }, + { + "expected_output": False, + "expected_error_message": "Input should be a valid string", + "expected_exception": ValidationError, + }, + ), + ( + { + "name": "Bookmark P", + "prompt": "Prompt A", + "prepared_prompt": "Prepared Prompt", + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": (), + "attack_module": "Module A", + "metric": "Metric A", + }, + { + "expected_output": False, + "expected_error_message": "Input should be a valid string", + "expected_exception": ValidationError, + }, + ), + ( + { + "name": "Bookmark P", + "prompt": "Prompt A", + "prepared_prompt": "Prepared Prompt", + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": [], + "attack_module": "Module A", + "metric": "Metric A", + }, + { + "expected_output": False, + "expected_error_message": "Input should be a valid string", + "expected_exception": ValidationError, + }, + ), + ( + { + "name": "Bookmark P", + "prompt": "Prompt A", + "prepared_prompt": "Prepared Prompt", + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": {}, + "attack_module": "Module A", + "metric": "Metric A", + }, + { + "expected_output": False, + "expected_error_message": "Input should be a valid string", + "expected_exception": ValidationError, + }, + ), + ( + { + "name": "Bookmark P", + "prompt": "Prompt A", + "prepared_prompt": "Prepared Prompt", + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": 123, + "attack_module": "Module A", + "metric": "Metric A", }, { "expected_output": False, @@ -390,11 +721,14 @@ def test_bookmark_singleton(self): # Invalid cases for 'attack_module' ( { - "name": "Bookmark Z", - "prompt": "Prompt Z", + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", - "response": "Response Z", - "attack_module": 123 + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": "Prompt Template", + "attack_module": None, + "metric": "Metric A", }, { "expected_output": False, @@ -404,11 +738,14 @@ def test_bookmark_singleton(self): ), ( { - "name": "Bookmark AA", - "prompt": "Prompt AA", - "response": "Response AA", + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", - "attack_module": [] + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": "Prompt Template", + "attack_module": (), + "metric": "Metric A", }, { "expected_output": False, @@ -418,11 +755,14 @@ def test_bookmark_singleton(self): ), ( { - "name": "Bookmark AB", - "prompt": "Prompt AB", - "response": "Response AB", + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", - "attack_module": {} + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": "Prompt Template", + "attack_module": [], + "metric": "Metric A", }, { "expected_output": False, @@ -430,14 +770,16 @@ def test_bookmark_singleton(self): "expected_exception": ValidationError, }, ), - # Invalid cases for 'context_strategy' ( { - "name": "Bookmark Z", - "prompt": "Prompt Z", - "response": "Response Z", + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", - "context_strategy": 123 + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": "Prompt Template", + "attack_module": {}, + "metric": "Metric A", }, { "expected_output": False, @@ -447,11 +789,14 @@ def test_bookmark_singleton(self): ), ( { - "name": "Bookmark AA", - "prompt": "Prompt AA", - "response": "Response AA", + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", - "context_strategy": [] + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": "Prompt Template", + "attack_module": 123, + "metric": "Metric A", }, { "expected_output": False, @@ -459,13 +804,17 @@ def test_bookmark_singleton(self): "expected_exception": ValidationError, }, ), + # Invalid cases for 'metric' ( { - "name": "Bookmark AB", - "prompt": "Prompt AB", - "response": "Response AB", + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", - "context_strategy": {} + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": "Prompt Template", + "attack_module": "Module A", + "metric": None, }, { "expected_output": False, @@ -473,14 +822,16 @@ def test_bookmark_singleton(self): "expected_exception": ValidationError, }, ), - # Invalid cases for 'prompt_template' ( { - "name": "Bookmark Z", - "prompt": "Prompt Z", - "response": "Response Z", + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", - "prompt_template": 123 + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": "Prompt Template", + "attack_module": "Module A", + "metric": (), }, { "expected_output": False, @@ -490,11 +841,14 @@ def test_bookmark_singleton(self): ), ( { - "name": "Bookmark AA", - "prompt": "Prompt AA", - "response": "Response AA", + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", - "prompt_template": [] + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": "Prompt Template", + "attack_module": "Module A", + "metric": [], }, { "expected_output": False, @@ -504,11 +858,14 @@ def test_bookmark_singleton(self): ), ( { - "name": "Bookmark AB", - "prompt": "Prompt AB", - "response": "Response AB", + "name": "Bookmark P", + "prompt": "Prompt A", "prepared_prompt": "Prepared Prompt", - "prompt_template": {} + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": "Prompt Template", + "attack_module": "Module A", + "metric": {}, }, { "expected_output": False, @@ -516,19 +873,41 @@ def test_bookmark_singleton(self): "expected_exception": ValidationError, }, ), - ] + ( + { + "name": "Bookmark P", + "prompt": "Prompt A", + "prepared_prompt": "Prepared Prompt", + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": "Prompt Template", + "attack_module": "Module A", + "metric": 123, + }, + { + "expected_output": False, + "expected_error_message": "Input should be a valid string", + "expected_exception": ValidationError, + }, + ), + ], ) def test_api_insert_bookmark(self, input_args, expected_dict): # Call the API function to insert a bookmark if expected_dict["expected_output"]: # Assert successful bookmark insertion result = api_insert_bookmark(**input_args) - assert result["success"] == expected_dict["expected_output"], "Bookmark insertion should succeed." + assert ( + result["success"] == expected_dict["expected_output"] + ), "Bookmark insertion should succeed." else: # Test invalid cases where bookmark insertion should fail and raise an exception with pytest.raises(expected_dict["expected_exception"]) as exc_info: api_insert_bookmark(**input_args) - + assert ( + exc_info.value.errors()[0]["msg"] + == expected_dict["expected_error_message"] + ) # ------------------------------------------------------------------------------ # Test api_get_all_bookmarks functionality @@ -537,119 +916,109 @@ def test_api_get_all_bookmark(self): # Set up the expected return value for get_all_bookmarks expected_bookmarks = [ { - 'name': 'Test Bookmark', - 'prompt': 'Test Prompt', + "name": "Test Bookmark", + "prompt": "Test Prompt", "prepared_prompt": "Test Prepared Prompt", - 'response': 'Test Response', - 'context_strategy': '', - 'prompt_template': '', - 'attack_module': '', - 'metric': '', - 'bookmark_time': '2024-07-14 22:26:51' + "response": "Test Response", + "context_strategy": "", + "prompt_template": "", + "attack_module": "", + "metric": "", + "bookmark_time": "2024-07-14 22:26:51", } ] # Call the API function actual_bookmarks = api_get_all_bookmarks() # Assert that the returned bookmarks match the expected bookmarks - assert actual_bookmarks == expected_bookmarks, "The returned bookmarks do not match the expected bookmarks." - - @pytest.mark.parametrize( - "bookmark_name, expected_dict", - [ - # Valid case - ( - "Test Bookmark", - { - "expected_result": { - 'name': 'Test Bookmark', - 'prompt': 'Test Prompt', - "prepared_prompt": "Test Prepared Prompt", - 'response': 'Test Response', - 'context_strategy': '', - 'prompt_template': '', - 'attack_module': '', - 'metric': '', - 'bookmark_time': '2024-07-14 22:26:51' - } - } - ), - ], -) - def test_api_get_bookmark(self, bookmark_name, expected_dict): - if "expected_exception" in expected_dict: - # Call the API function and assert that the expected exception is raised - with pytest.raises(expected_dict["expected_exception"]) as exc_info: - api_get_bookmark(bookmark_name) - # Assert that the exception message matches the expected message - assert str(exc_info.value) == expected_dict["expected_message"], "The expected exception message was not raised." - else: - # Call the API function and get the response - response = api_get_bookmark(bookmark_name) - # Assert that the response matches the expected bookmark - assert response == expected_dict["expected_result"], "The returned bookmark does not match the expected bookmark." - + assert ( + actual_bookmarks == expected_bookmarks + ), "The returned bookmarks do not match the expected bookmarks." # ------------------------------------------------------------------------------ - # Test api_get_delete_all_bookmark functionality + # Test api_get_bookmark functionality # ------------------------------------------------------------------------------ - def test_api_delete_all_bookmark(self): - # Set up the expected return value for delete_all_bookmark - expected_response = {'success': True, 'message': 'All bookmark records deleted.'} - - # Call the API function - delete_response = api_delete_all_bookmark() - - # Assert that the response matches the expected response - assert delete_response == expected_response, "The response from delete_all_bookmark does not match the expected response." - + @pytest.mark.parametrize( + "bookmark_name, expected_result", + [ + # Valid case + ( + "Test Bookmark", + { + "name": "Test Bookmark", + "prompt": "Test Prompt", + "prepared_prompt": "Test Prepared Prompt", + "response": "Test Response", + "context_strategy": "", + "prompt_template": "", + "attack_module": "", + "metric": "", + "bookmark_time": "2024-07-14 22:26:51", + }, + ), + ], + ) + def test_api_get_bookmark(self, bookmark_name, expected_result, mocker): + # Mock the get_bookmark method of the Bookmark class + mocker.patch.object(Bookmark, 'get_bookmark', return_value=expected_result) + + # Call the API function and get the result + result = api_get_bookmark(bookmark_name) + + # Assert that the result matches the expected result + assert result == expected_result, "The returned result does not match the expected result." # ------------------------------------------------------------------------------ # Test api_delete_bookmark functionality # ------------------------------------------------------------------------------ @pytest.mark.parametrize( - "bookmark_name, expected_dict", + "bookmark_name, expected_result", [ # Valid case ( "Test Bookmark", { - "expected_result": { - 'success': True, 'message': 'Bookmark record deleted.' - } - } - ) + "success": True, + "message": "[Bookmark] Bookmark record deleted.", + }, + ), ], ) - def test_api_delete_bookmark(self, bookmark_name, expected_dict): - # Extract variables from expected_dict - expected_exception = expected_dict.get("expected_exception") - expected_message = expected_dict.get("expected_message") - expected_result = expected_dict.get("expected_result") + def test_api_delete_bookmark(self, bookmark_name, expected_result): + # Call the API function and get the response + response = api_delete_bookmark(bookmark_name) + # Assert that the response matches the expected result + assert response == expected_result, "The returned result does not match the expected result." - if expected_exception: - # Call the API function and assert that the expected exception is raised - with pytest.raises(expected_exception) as exc_info: - api_delete_bookmark(bookmark_name) - # Assert that the exception message matches the expected message - assert str(exc_info.value) == expected_message, "The expected exception message was not raised." - else: - # Call the API function and get the response - response = api_delete_bookmark(bookmark_name) - # Assert that the response matches the expected result - assert response == expected_result, "The returned result does not match the expected result." + # ------------------------------------------------------------------------------ + # Test api_get_delete_all_bookmark functionality + # ------------------------------------------------------------------------------ + def test_api_delete_all_bookmark(self): + # Set up the expected return value for delete_all_bookmark + expected_response = { + "success": True, + "message": "[Bookmark] All bookmark records deleted.", + } + + # Call the API function + delete_response = api_delete_all_bookmark() + + # Assert that the response matches the expected response + assert ( + delete_response == expected_response + ), "The response from delete_all_bookmark does not match the expected response." # ------------------------------------------------------------------------------ # Test api_export_bookmark functionality # ------------------------------------------------------------------------------ @pytest.mark.parametrize( "export_file_name, expected_output", - [ - ("bookmark", "tests/unit-tests/src/data/bookmarks/bookmark.json") - ] + [("bookmark", "tests/unit-tests/src/data/bookmarks/bookmark.json")], ) def test_api_export_bookmarks(self, export_file_name, expected_output): # Call the API function actual_output = api_export_bookmarks(export_file_name) # Assert that the returned output matches the expected output - assert actual_output == expected_output, "The returned output from api_export_bookmarks does not match the expected output." + assert ( + actual_output == expected_output + ), "The returned output from api_export_bookmarks does not match the expected output." diff --git a/tests/unit-tests/src/test_bookmark.py b/tests/unit-tests/src/test_bookmark.py new file mode 100644 index 00000000..dba9ee8c --- /dev/null +++ b/tests/unit-tests/src/test_bookmark.py @@ -0,0 +1,913 @@ +import textwrap +from unittest.mock import patch + +import pytest + +from moonshot.src.bookmark.bookmark import Bookmark +from moonshot.src.bookmark.bookmark_arguments import BookmarkArguments + + +class TestCollectionBookmark: + # ------------------------------------------------------------------------------ + # Test add_bookmark functionality + # ------------------------------------------------------------------------------ + @pytest.mark.parametrize( + "input_args, mock_return_value, expected_dict", + [ + # Valid case for bookmark insertion + ( + { + "name": "Bookmark A", + "prompt": "Prompt A", + "prepared_prompt": "Prepared Prompt", + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", + "bookmark_time": "Time A", + }, + True, + { + "expected_output": True, + "expected_message": "[Bookmark] Bookmark added successfully.", + }, + ), + # Case where insertion fails + ( + { + "name": "Bookmark B", + "prompt": "Prompt B", + "prepared_prompt": "Prepared Prompt", + "response": "Response B", + "context_strategy": "Strategy B", + "prompt_template": "Template B", + "attack_module": "Module B", + "metric": "Metric B", + "bookmark_time": "Time B", + }, + None, + { + "expected_output": False, + "expected_message": "[Bookmark] Failed to add bookmark record: Error " + "inserting record into database.", + }, + ), + ], + ) + @patch("moonshot.src.storage.storage.Storage.create_database_record") + def test_add_bookmark( + self, mock_create_database_record, input_args, mock_return_value, expected_dict + ): + bookmark_instance = Bookmark() + + # Convert input_args to BookmarkArguments + bookmark_args = BookmarkArguments(**input_args) + + # Set up the mock return value for create_database_record + mock_create_database_record.return_value = mock_return_value + + # Call the API function to insert a bookmark + result = bookmark_instance.add_bookmark(bookmark_args) + + # Assert the result + assert ( + result["success"] == expected_dict["expected_output"] + ), "Bookmark insertion success status should match." + assert ( + result["message"] == expected_dict["expected_message"] + ), "Bookmark insertion message should match." + + # Close the instance + bookmark_instance.close() + + # ------------------------------------------------------------------------------ + # Test get_all_bookmarks functionality + # ------------------------------------------------------------------------------ + @pytest.mark.parametrize( + "mock_return_value, expected_output", + [ + # Case with multiple bookmarks + ( + [ + ( + "1", + "Bookmark A", + "Prompt A", + "Prepared Prompt A", + "Response A", + "Strategy A", + "Template A", + "Module A", + "Metric A", + "Time A", + ), + ( + "2", + "Bookmark B", + "Prompt B", + "Prepared Prompt B", + "Response B", + "Strategy B", + "Template B", + "Module B", + "Metric B", + "Time B", + ), + ], + [ + { + "name": "Bookmark A", + "prompt": "Prompt A", + "prepared_prompt": "Prepared Prompt A", + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", + "bookmark_time": "Time A", + }, + { + "name": "Bookmark B", + "prompt": "Prompt B", + "prepared_prompt": "Prepared Prompt B", + "response": "Response B", + "context_strategy": "Strategy B", + "prompt_template": "Template B", + "attack_module": "Module B", + "metric": "Metric B", + "bookmark_time": "Time B", + }, + ], + ), + # Case with no bookmarks + ( + [], + [], + ), + # Case where read_database_records returns None + ( + None, + [], + ), + # Case where read_database_records returns an empty string + ( + "", + [], + ), + # Case where read_database_records returns a dictionary + ( + {"key": "value"}, + [], + ), + # Case where read_database_records returns a list of strings + ( + ["string1", "string2"], + [], + ), + # Case where read_database_records returns an integer + ( + 123, + [], + ), + ], + ) + @patch("moonshot.src.storage.storage.Storage.read_database_records") + def test_get_all_bookmarks( + self, mock_read_database_records, mock_return_value, expected_output + ): + bookmark_instance = Bookmark() + + # Set up the mock return value for read_database_records + mock_read_database_records.return_value = mock_return_value + + # Call the API function to retrieve all bookmarks + result = bookmark_instance.get_all_bookmarks() + + # Assert the result + assert ( + result == expected_output + ), "The list of bookmarks should match the expected output." + + # Assert read_database_records was called once + mock_read_database_records.assert_called_once_with( + bookmark_instance.db_instance, + Bookmark.sql_select_bookmarks_record, + ) + + # Close the instance + bookmark_instance.close() + + # ------------------------------------------------------------------------------ + # Test get_bookmark functionality + # ------------------------------------------------------------------------------ + @pytest.mark.parametrize( + "bookmark_name, mock_return_value, expected_output, expected_exception, expected_call", + [ + # Case where the bookmark is found + ( + "Bookmark A", + ( + "1", + "Bookmark A", + "Prompt A", + "Prepared Prompt A", + "Response A", + "Strategy A", + "Template A", + "Module A", + "Metric A", + "Time A", + ), + { + "name": "Bookmark A", + "prompt": "Prompt A", + "prepared_prompt": "Prepared Prompt A", + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", + "bookmark_time": "Time A", + }, + None, + True, + ), + # Case where the bookmark_name is invalid + ( + None, + ( + "1", + "Bookmark A", + "Prompt A", + "Prepared Prompt A", + "Response A", + "Strategy A", + "Template A", + "Module A", + "Metric A", + "Time A", + ), + "[Bookmark] Invalid bookmark name: None", + RuntimeError, + False, + ), + ( + "", + ( + "1", + "Bookmark A", + "Prompt A", + "Prepared Prompt A", + "Response A", + "Strategy A", + "Template A", + "Module A", + "Metric A", + "Time A", + ), + "[Bookmark] Invalid bookmark name: ", + RuntimeError, + False, + ), + ( + (), + ( + "1", + "Bookmark A", + "Prompt A", + "Prepared Prompt A", + "Response A", + "Strategy A", + "Template A", + "Module A", + "Metric A", + "Time A", + ), + "[Bookmark] Invalid bookmark name: ()", + RuntimeError, + False, + ), + ( + [], + ( + "1", + "Bookmark A", + "Prompt A", + "Prepared Prompt A", + "Response A", + "Strategy A", + "Template A", + "Module A", + "Metric A", + "Time A", + ), + "[Bookmark] Invalid bookmark name: []", + RuntimeError, + False, + ), + ( + {}, + ( + "1", + "Bookmark A", + "Prompt A", + "Prepared Prompt A", + "Response A", + "Strategy A", + "Template A", + "Module A", + "Metric A", + "Time A", + ), + "[Bookmark] Invalid bookmark name: {}", + RuntimeError, + False, + ), + ( + 123, + ( + "1", + "Bookmark A", + "Prompt A", + "Prepared Prompt A", + "Response A", + "Strategy A", + "Template A", + "Module A", + "Metric A", + "Time A", + ), + "[Bookmark] Invalid bookmark name: 123", + RuntimeError, + False, + ), + # Case where read_database_record returns an unexpected type (e.g., a dictionary) + ( + "Bookmark C", + None, + "[Bookmark] No record found for bookmark name: Bookmark C", + RuntimeError, + True, + ), + ( + "Bookmark C", + "", + "[Bookmark] No record found for bookmark name: Bookmark C", + RuntimeError, + True, + ), + ( + "Bookmark C", + [], + "[Bookmark] No record found for bookmark name: Bookmark C", + RuntimeError, + True, + ), + ( + "Bookmark C", + {}, + "[Bookmark] No record found for bookmark name: Bookmark C", + RuntimeError, + True, + ), + ( + "Bookmark C", + 123, + "[Bookmark] No record found for bookmark name: Bookmark C", + RuntimeError, + True, + ), + ], + ) + @patch("moonshot.src.storage.storage.Storage.read_database_record") + def test_get_bookmark( + self, + mock_read_database_record, + bookmark_name, + mock_return_value, + expected_output, + expected_exception, + expected_call, + ): + bookmark_instance = Bookmark() + + # Set up the mock return value for read_database_record + mock_read_database_record.return_value = mock_return_value + + # Call the API function to retrieve the bookmark + if expected_exception: + with pytest.raises(expected_exception): + bookmark_instance.get_bookmark(bookmark_name) + else: + result = bookmark_instance.get_bookmark(bookmark_name) + # Assert the result + assert ( + result == expected_output + ), "The bookmark information should match the expected output." + + # Assert read_database_record was called with the correct arguments + if expected_call: + mock_read_database_record.assert_called_once_with( + bookmark_instance.db_instance, + (bookmark_name,), + Bookmark.sql_select_bookmark_record, + ) + else: + mock_read_database_record.assert_not_called() + + # Close the instance + bookmark_instance.close() + + # ------------------------------------------------------------------------------ + # Test delete_bookmark functionality + # ------------------------------------------------------------------------------ + @pytest.mark.parametrize( + "mock_side_effect, bookmark_name, expected_output, expected_call", + [ + # Case where the bookmark is deleted successfully + ( + None, + "Bookmark A", + {"success": True, "message": "[Bookmark] Bookmark record deleted."}, + True, + ), + # Case where the bookmark_name is invalid (None) + ( + None, + None, + {"success": False, "message": "[Bookmark] Invalid bookmark name: None"}, + False, + ), + # Case where the bookmark_name is invalid (empty string) + ( + None, + "", + {"success": False, "message": "[Bookmark] Invalid bookmark name: "}, + False, + ), + # Case where the bookmark_name is invalid (tuple) + ( + None, + (), + {"success": False, "message": "[Bookmark] Invalid bookmark name: ()"}, + False, + ), + # Case where the bookmark_name is invalid (list) + ( + None, + [], + {"success": False, "message": "[Bookmark] Invalid bookmark name: []"}, + False, + ), + # Case where the bookmark_name is invalid (dictionary) + ( + None, + {}, + {"success": False, "message": "[Bookmark] Invalid bookmark name: {}"}, + False, + ), + # Case where the bookmark_name is invalid (integer) + ( + None, + 123, + {"success": False, "message": "[Bookmark] Invalid bookmark name: 123"}, + False, + ), + # Case where deletion fails + ( + Exception("Deletion error"), + "Bookmark A", + { + "success": False, + "message": "[Bookmark] Failed to delete bookmark record: Deletion error", + }, + True, + ), + ], + ) + @patch("moonshot.src.storage.storage.Storage.delete_database_record_in_table") + def test_delete_bookmark( + self, + mock_delete_database_record_in_table, + mock_side_effect, + bookmark_name, + expected_output, + expected_call, + ): + bookmark_instance = Bookmark() + + # Set up the mock side effect for delete_database_record_in_table + mock_delete_database_record_in_table.side_effect = mock_side_effect + + # Call the API function to delete the bookmark + result = bookmark_instance.delete_bookmark(bookmark_name) + + # Assert the result + assert ( + result == expected_output + ), "The result of deleting the bookmark should match the expected output." + + # Assert delete_database_record_in_table was called with the correct arguments + if expected_call: + sql_delete_bookmark_record = textwrap.dedent( + f""" + DELETE FROM bookmark WHERE name = '{bookmark_name}'; + """ + ) + mock_delete_database_record_in_table.assert_called_once_with( + bookmark_instance.db_instance, + sql_delete_bookmark_record, + ) + else: + mock_delete_database_record_in_table.assert_not_called() + + # Close the instance + bookmark_instance.close() + + # ------------------------------------------------------------------------------ + # Test delete_all_bookmark functionality + # ------------------------------------------------------------------------------ + @pytest.mark.parametrize( + "mock_side_effect, expected_output", + [ + # Case where deletion is successful + ( + None, + { + "success": True, + "message": "[Bookmark] All bookmark records deleted.", + }, + ), + # Case where deletion fails + ( + Exception("Deletion error"), + { + "success": False, + "message": "[Bookmark] Failed to delete all bookmark records: Deletion error", + }, + ), + ], + ) + @patch("moonshot.src.storage.storage.Storage.delete_database_record_in_table") + def test_delete_all_bookmark( + self, mock_delete_database_record_in_table, mock_side_effect, expected_output + ): + bookmark_instance = Bookmark() + + # Set up the mock side effect for delete_database_record_in_table + mock_delete_database_record_in_table.side_effect = mock_side_effect + + # Call the API function to delete all bookmarks + result = bookmark_instance.delete_all_bookmark() + + # Assert the result + assert ( + result == expected_output + ), "The result of deleting all bookmarks should match the expected output." + + # Assert delete_database_record_in_table was called once with the correct arguments + mock_delete_database_record_in_table.assert_called_once_with( + bookmark_instance.db_instance, + Bookmark.sql_delete_bookmark_records, + ) + + # Close the instance + bookmark_instance.close() + + # ------------------------------------------------------------------------------ + # Test export_bookmarks functionality + # ------------------------------------------------------------------------------ + @pytest.mark.parametrize( + "mock_read_return_value, mock_create_return_value, expected_bookmarks_json, expected_output", + [ + # Case where export is successful with multiple bookmarks + ( + [ + ( + "1", + "Bookmark A", + "Prompt A", + "Prepared Prompt A", + "Response A", + "Strategy A", + "Template A", + "Module A", + "Metric A", + "Time A", + ), + ( + "2", + "Bookmark B", + "Prompt B", + "Prepared Prompt B", + "Response B", + "Strategy B", + "Template B", + "Module B", + "Metric B", + "Time B", + ), + ], + "moonshot-data/bookmark/bookmarks.json", + [ + { + "name": "Bookmark A", + "prompt": "Prompt A", + "prepared_prompt": "Prepared Prompt A", + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", + "bookmark_time": "Time A", + }, + { + "name": "Bookmark B", + "prompt": "Prompt B", + "prepared_prompt": "Prepared Prompt B", + "response": "Response B", + "context_strategy": "Strategy B", + "prompt_template": "Template B", + "attack_module": "Module B", + "metric": "Metric B", + "bookmark_time": "Time B", + }, + ], + "moonshot-data/bookmark/bookmarks.json", + ), + # Case where export is successful with no bookmarks + ( + [], + "moonshot-data/bookmark/bookmarks.json", + [], + "moonshot-data/bookmark/bookmarks.json", + ), + # Case where read_database_records returns None + ( + None, + "moonshot-data/bookmark/bookmarks.json", + [], + "moonshot-data/bookmark/bookmarks.json", + ), + # Case where read_database_records returns an empty string + ( + "", + "moonshot-data/bookmark/bookmarks.json", + [], + "moonshot-data/bookmark/bookmarks.json", + ), + # Case where read_database_records returns a dictionary + ( + {"key": "value"}, + "moonshot-data/bookmark/bookmarks.json", + [], + "moonshot-data/bookmark/bookmarks.json", + ), + # Case where read_database_records returns a list of strings + ( + ["string1", "string2"], + "moonshot-data/bookmark/bookmarks.json", + [], + "moonshot-data/bookmark/bookmarks.json", + ), + # Case where read_database_records returns an integer + ( + 123, + "moonshot-data/bookmark/bookmarks.json", + [], + "moonshot-data/bookmark/bookmarks.json", + ), + ], + ) + @patch("moonshot.src.storage.storage.Storage.read_database_records") + @patch("moonshot.src.storage.storage.Storage.create_object") + def test_export_bookmarks( + self, + mock_create_object, + mock_read_database_records, + mock_read_return_value, + mock_create_return_value, + expected_bookmarks_json, + expected_output, + ): + bookmark_instance = Bookmark() + + # Set up the mock return values + mock_read_database_records.return_value = mock_read_return_value + mock_create_object.return_value = mock_create_return_value + + # Call the API function to export bookmarks + result = bookmark_instance.export_bookmarks() + + # Assert the result + assert ( + result == expected_output + ), "The path to the exported JSON file should match the expected output." + + # Assert read_database_records was called once with the correct arguments + mock_read_database_records.assert_called_once_with( + bookmark_instance.db_instance, + Bookmark.sql_select_bookmarks_record, + ) + + # Assert create_object was called once with the correct arguments + mock_create_object.assert_called_once_with( + "BOOKMARKS", + "bookmarks", + {"bookmarks": expected_bookmarks_json}, + "json", + ) + + # Close the instance + bookmark_instance.close() + + @pytest.mark.parametrize( + "export_file_name, expected_exception, expected_message", + [ + # Case where export_file_name is None + ( + None, + Exception, + "[Bookmark] Failed to export bookmarks: Export filename must be a non-empty string", + ), + # Case where export_file_name is an empty string + ( + "", + Exception, + "[Bookmark] Failed to export bookmarks: Export filename must be a non-empty string", + ), + # Case where export_file_name is a tuple + ( + (), + Exception, + "[Bookmark] Failed to export bookmarks: Export filename must be a non-empty string", + ), + # Case where export_file_name is a list + ( + [], + Exception, + "[Bookmark] Failed to export bookmarks: Export filename must be a non-empty string", + ), + # Case where export_file_name is a dictionary + ( + {}, + Exception, + "[Bookmark] Failed to export bookmarks: Export filename must be a non-empty string", + ), + # Case where export_file_name is an integer + ( + 123, + Exception, + "[Bookmark] Failed to export bookmarks: Export filename must be a non-empty string", + ), + ], + ) + def test_export_bookmarks_invalid_filename( + self, export_file_name, expected_exception, expected_message + ): + bookmark_instance = Bookmark() + + # Call the API function to export bookmarks and assert the exception + with pytest.raises(expected_exception) as exc_info: + bookmark_instance.export_bookmarks(export_file_name) + + # Assert the exception message + assert ( + str(exc_info.value) == expected_message + ), "The exception message should match the expected message." + + @pytest.mark.parametrize( + "mock_read_return_value, mock_create_side_effect, expected_exception, expected_message", + [ + # Case where create_object raises an exception + ( + [ + ( + "1", + "Bookmark A", + "Prompt A", + "Prepared Prompt A", + "Response A", + "Strategy A", + "Template A", + "Module A", + "Metric A", + "Time A", + ), + ], + Exception("Export error"), + Exception, + "[Bookmark] Failed to export bookmarks: Export error", + ), + ], + ) + @patch("moonshot.src.storage.storage.Storage.read_database_records") + @patch("moonshot.src.storage.storage.Storage.create_object") + def test_export_bookmarks_exception( + self, + mock_create_object, + mock_read_database_records, + mock_read_return_value, + mock_create_side_effect, + expected_exception, + expected_message, + ): + bookmark_instance = Bookmark() + + # Set up the mock return values and side effects + mock_read_database_records.return_value = mock_read_return_value + mock_create_object.side_effect = mock_create_side_effect + + # Call the API function to export bookmarks and assert the exception + with pytest.raises(expected_exception) as exc_info: + bookmark_instance.export_bookmarks() + + # Assert the exception message + assert ( + str(exc_info.value) == expected_message + ), "The exception message should match the expected message." + + # Assert read_database_records was called once with the correct arguments + mock_read_database_records.assert_called_once_with( + bookmark_instance.db_instance, + Bookmark.sql_select_bookmarks_record, + ) + + # Close the instance + bookmark_instance.close() + + # ------------------------------------------------------------------------------ + # Test close functionality + # ------------------------------------------------------------------------------ + @patch("moonshot.src.storage.storage.Storage.close_database_connection") + def test_close(self, mock_close_database_connection): + # Create an instance of Bookmark + bookmark_instance = Bookmark() + + # Call the close method + bookmark_instance.close() + + # Assert close_database_connection was called once with the correct arguments + mock_close_database_connection.assert_called_once_with( + bookmark_instance.db_instance + ) + + # Assert the singleton instance is set to None + assert ( + Bookmark._instance is None + ), "The singleton instance should be set to None after closing." + + @patch("moonshot.src.storage.storage.Storage.close_database_connection") + def test_close_no_db_instance(self, mock_close_database_connection): + # Create an instance of Bookmark + bookmark_instance = Bookmark() + bookmark_instance.db_instance = None # Simulate no database instance + + # Call the close method + bookmark_instance.close() + + # Assert close_database_connection was not called + mock_close_database_connection.assert_not_called() + + # Assert the singleton instance is set to None + assert ( + Bookmark._instance is None + ), "The singleton instance should be set to None after closing." + + @patch("moonshot.src.storage.storage.Storage.close_database_connection") + def test_close_multiple_times(self, mock_close_database_connection): + # Create an instance of Bookmark + bookmark_instance = Bookmark() + + # Call the close method multiple times + bookmark_instance.close() + bookmark_instance.close() + bookmark_instance.close() + + # Assert close_database_connection was called three times with the correct arguments + assert ( + mock_close_database_connection.call_count == 3 + ), "close_database_connection should be called three times." + mock_close_database_connection.assert_called_with(bookmark_instance.db_instance) + + # Assert the singleton instance is set to None + assert ( + Bookmark._instance is None + ), "The singleton instance should be set to None after closing." + + # ------------------------------------------------------------------------------ + # Test bookmark instance is singleton + # ------------------------------------------------------------------------------ + def test_bookmark_singleton(self): + # Retrieve two instances of the Bookmark class + bookmark_instance_1 = Bookmark() + bookmark_instance_2 = Bookmark() + + # Assert that both instances are the same (singleton behavior) + assert ( + bookmark_instance_1 is bookmark_instance_2 + ), "Bookmark instances should be the same (singleton pattern)." + + # Terminate the bookmark_instance + bookmark_instance_1.close() diff --git a/tests/unit-tests/src/test_bookmark_arguments.py b/tests/unit-tests/src/test_bookmark_arguments.py new file mode 100644 index 00000000..e7f5447c --- /dev/null +++ b/tests/unit-tests/src/test_bookmark_arguments.py @@ -0,0 +1,100 @@ +import pytest +from pydantic import ValidationError + +from moonshot.src.bookmark.bookmark_arguments import BookmarkArguments + + +class TestCollectionBookmarkArguments: + def test_create_bookmark_arguments(self): + # Test creating a valid BookmarkArguments instance + bookmark_args = BookmarkArguments( + name="Bookmark A", + prompt="Prompt A", + prepared_prompt="Prepared Prompt A", + response="Response A", + context_strategy="Strategy A", + prompt_template="Template A", + attack_module="Module A", + metric="Metric A", + bookmark_time="Time A", + ) + assert bookmark_args.name == "Bookmark A" + assert bookmark_args.prompt == "Prompt A" + assert bookmark_args.prepared_prompt == "Prepared Prompt A" + assert bookmark_args.response == "Response A" + assert bookmark_args.context_strategy == "Strategy A" + assert bookmark_args.prompt_template == "Template A" + assert bookmark_args.attack_module == "Module A" + assert bookmark_args.metric == "Metric A" + assert bookmark_args.bookmark_time == "Time A" + + @pytest.mark.parametrize( + "name, prompt, prepared_prompt, response", + [ + ("", "Prompt A", "Prepared Prompt A", "Response A"), # Invalid name + ("Bookmark A", "", "Prepared Prompt A", "Response A"), # Invalid prompt + ("Bookmark A", "Prompt A", "", "Response A"), # Invalid prepared_prompt + ("Bookmark A", "Prompt A", "Prepared Prompt A", ""), # Invalid response + ], + ) + def test_create_bookmark_arguments_invalid( + self, name, prompt, prepared_prompt, response + ): + with pytest.raises(ValidationError): + BookmarkArguments( + name=name, + prompt=prompt, + prepared_prompt=prepared_prompt, + response=response, + context_strategy="Strategy A", + prompt_template="Template A", + attack_module="Module A", + metric="Metric A", + bookmark_time="Time A", + ) + + def test_from_tuple_to_dict(self): + # Test converting a tuple to a dictionary + values = ( + "1", + "Bookmark A", + "Prompt A", + "Prepared Prompt A", + "Response A", + "Strategy A", + "Template A", + "Module A", + "Metric A", + "Time A", + ) + expected_dict = { + "name": "Bookmark A", + "prompt": "Prompt A", + "prepared_prompt": "Prepared Prompt A", + "response": "Response A", + "context_strategy": "Strategy A", + "prompt_template": "Template A", + "attack_module": "Module A", + "metric": "Metric A", + "bookmark_time": "Time A", + } + result_dict = BookmarkArguments.from_tuple_to_dict(values) + assert ( + result_dict == expected_dict + ), "The dictionary representation should match the expected output." + + def test_from_tuple_to_dict_insufficient_values(self): + # Test converting a tuple with insufficient values to a dictionary + values = ( + "1", + "Bookmark A", + "Prompt A", + "Prepared Prompt A", + "Response A", + ) + with pytest.raises(ValueError) as exc: + BookmarkArguments.from_tuple_to_dict(values) + assert ( + str(exc.value) + == "[BookmarkArguments] Failed to convert to dictionary because of the insufficient number of values" + ), "The error message should match the expected validation error message."