Update run-tests.py

Refactor Test Runner Script for Robustness and Efficiency

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:])
+    # 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)