diff --git a/tests/run-tests.py b/tests/run-tests.py index 0277ccd3..263538c8 100755 --- a/tests/run-tests.py +++ b/tests/run-tests.py @@ -1,66 +1,102 @@ #!/usr/bin/env python import os import pytest -import shutil +import shutil import subprocess import sys from utils import Logger -class RunTests(): - # Grab the folders for the llmware and llmware-packaging repos. They are assumed to live in the same parent folder +class RunTests: + # Initialize and grab the folder paths for the llmware and llmware-packaging repos def __init__(self): tests_folder = os.path.dirname(os.path.realpath(__file__)) - self.repo_root = os.path.join(tests_folder, "..") - - # Ensure the latest/greatest llmware module is installed locally + self.repo_root = os.path.abspath(os.path.join(tests_folder, "..")) + self.logger = Logger() + + # Ensure the latest llmware module is installed locally def update_llmware_install(self): - self.run_command("pip uninstall llmware -y", self.repo_root) - self.run_command("pip install .", self.repo_root) + self.logger.log("Updating llmware installation...") + try: + self.run_command("pip uninstall llmware -y", self.repo_root) + self.run_command("pip install .", self.repo_root) + except Exception as e: + self.logger.log(f"Error updating llmware: {e}", level="error") + sys.exit(1) - # Try to start from as clean an environment as possible + # Clean the environment to start from a fresh state def clean_the_environment(self): - Logger().log("Cleaning the environment...") + self.logger.log("Cleaning the environment...") + + # Remove the default llmware data folders + self.remove_folder(os.path.join(os.environ["HOME"], "llmware_data")) + self.remove_folder(os.path.join(os.environ["HOME"], "llmware_data_custom")) - # Remove the default data folder: $HOME/llmware_data - llmware_data_dir = os.path.join(os.environ["HOME"],"llmware_data") - llmware_data_custom_dir = os.path.join(os.environ["HOME"],"llmware_data_custom") - Logger().log(f"Deleting {llmware_data_dir}") - if os.path.exists(llmware_data_dir): - shutil.rmtree(llmware_data_dir) + # Reset MongoDB + try: + self.logger.log("Resetting MongoDB (dropping the 'llmware' database)...") + from llmware.resources import MongoDBManager + MongoDBManager().client.drop_database("llmware") + except Exception as e: + self.logger.log(f"Error resetting MongoDB: {e}", level="error") + sys.exit(1) - if os.path.exists(llmware_data_custom_dir): - shutil.rmtree(llmware_data_custom_dir) - - # Reset Mongo - Logger().log("Resetting Mongo (deleting the llmware db)") - from llmware.resources import MongoDBManager - MongoDBManager().client.drop_database("llmware") + # Reset Milvus collections + try: + self.logger.log("Resetting Milvus (dropping all collections)...") + from llmware.configs import LLMWareConfig + from pymilvus import connections, utility + connections.connect("default", host=os.environ.get('MILVUS_HOST', 'localhost'), port=19530) + for collection in utility.list_collections(): + utility.drop_collection(collection) + self.logger.log("All Milvus collections dropped successfully.") + except Exception as e: + self.logger.log(f"Error resetting Milvus: {e}", level="error") + sys.exit(1) - # Reset Milvus - Logger().log("Resetting Milvus (deleting all collections)") - from llmware.configs import LLMWareConfig - from pymilvus import connections, utility, FieldSchema, CollectionSchema, DataType, Collection - connections.connect("default", host=os.environ.get('MILVUS_HOST','localhost') , port=19530) - for collection in utility.list_collections(): - utility.drop_collection(collection) + # Helper function to remove a folder if it exists + def remove_folder(self, folder_path): + if os.path.exists(folder_path): + try: + self.logger.log(f"Removing folder: {folder_path}") + shutil.rmtree(folder_path) + except Exception as e: + self.logger.log(f"Error removing folder {folder_path}: {e}", level="error") + else: + self.logger.log(f"Folder not found: {folder_path}") - # Run a system command in the given dir, wait for it to complete and print the output + # Run a system command in the given directory, wait for it to complete, and log the output def run_command(self, command, working_dir): - Logger().log(f"Running command '{command}' in '{working_dir}'...") + self.logger.log(f"Running command '{command}' in '{working_dir}'...") command_array = command.split(" ") - command_output = "" - p = subprocess.Popen(command_array, cwd=working_dir, stdout=subprocess.PIPE) - for line in p.stdout: - output_line = line.decode('utf8') - command_output += output_line - Logger().log(output_line) - p.wait() - Logger().log("") - return command_output + try: + process = subprocess.Popen(command_array, cwd=working_dir, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + stdout, stderr = process.communicate() + if stdout: + self.logger.log(stdout.decode('utf-8')) + if stderr: + self.logger.log(stderr.decode('utf-8'), level="error") + if process.returncode != 0: + raise subprocess.CalledProcessError(process.returncode, command) + except subprocess.CalledProcessError as e: + self.logger.log(f"Command '{command}' failed with error: {e}", level="error") + sys.exit(1) + except Exception as e: + self.logger.log(f"An unexpected error occurred while running '{command}': {e}", level="error") + sys.exit(1) + + +if __name__ == "__main__": + test_runner = RunTests() + + # Update and clean environment before running tests + test_runner.update_llmware_install() + test_runner.clean_the_environment() -test_runner = RunTests() -test_runner.update_llmware_install() -test_runner.clean_the_environment() -pytest.main(sys.argv[1:]) \ No newline at end of file + # Run the tests with pytest + try: + pytest.main(sys.argv[1:]) + except Exception as e: + test_runner.logger.log(f"Error running tests: {e}", level="error") + sys.exit(1)