From 8e487863cf94aeef357d207e67fa76a8406d0a81 Mon Sep 17 00:00:00 2001
From: VIRTUALGUARD <virtualguard101@gmail.com>
Date: Fri, 11 Jul 2025 14:06:19 +0800
Subject: [PATCH] Update test script

---
 test/main.py | 72 +++++++++++++++++++++++++++++++++-------------------
 1 file changed, 46 insertions(+), 26 deletions(-)

diff --git a/test/main.py b/test/main.py
index c6bf2ee..5ad025b 100644
--- a/test/main.py
+++ b/test/main.py
@@ -1,25 +1,35 @@
+#!/usr/bin/env python3
+
 from src.manager import *
 import logging
 import colorlog
+import sys
 
-# Initialize color logging
-handler = colorlog.StreamHandler()
-handler.setFormatter(colorlog.ColoredFormatter(
-    '%(log_color)s%(asctime)s - %(levelname)s - %(message)s',
-    log_colors={
-        'DEBUG':    'cyan',
-        'INFO':     'green',
-        'WARNING':  'yellow',
-        'ERROR':    'red',
-        'CRITICAL': 'bold_red',
-    }
-))
+def initialize_logger():
+    """Initialize and return a color logger."""
+    handler = colorlog.StreamHandler()
+    handler.setFormatter(colorlog.ColoredFormatter(
+        '%(log_color)s%(asctime)s - %(levelname)s - %(message)s',
+        log_colors={
+            'DEBUG':    'cyan',
+            'INFO':     'green',
+            'WARNING':  'yellow',
+            'ERROR':    'red',
+            'CRITICAL': 'bold_red',
+        }
+    ))
 
-logger = colorlog.getLogger(__name__)
-logger.setLevel(logging.INFO)
-logger.addHandler(handler)
+    logger = colorlog.getLogger(__name__)
+    logger.setLevel(logging.INFO)
+    logger.addHandler(handler)
+    return logger
 
-service_manager = Manager()
+def initialize_service_manager():
+    """Initialize and return the service manager."""
+    return Manager()
+
+logger = initialize_logger()
+service_manager = initialize_service_manager()
 
 def start_menu():
     logger.info(f"=======================================")
@@ -31,23 +41,26 @@ def start_menu():
     logger.info(f"=====> c. Count the number of services")
     logger.info(f"=====> q. Quit ")
 
-def main():
+def main_menu():
+    """Display the main menu and handle user input."""
+
     while True:
         start_menu()
 
-        start_choice = input()
+        start_choice = input("Enter your choice: ").strip()
         if start_choice == '1':
             service_tag = input("Please input the deploy way of service you want to register('docker' | 'sys'): ")
             service_name = input("Please input the name of the service: ")
             if service_tag == "docker":
                 service_path = input("Please input the config path of the service where 'docker-compose.yml' located: ")
                 service_manager.register_service(service_tag, service_name, service_path)
-            service_manager.register_service(service_tag, service_name)
+            else:
+                service_manager.register_service(service_tag, service_name)
         if start_choice == '2':
-            if len(service_manager.services_list) == 0:
-                raise ValueError("There has no service registered")
+            if not service_manager.services_list:
+                logger.error("No services registered.")
+                continue
             while True:
-                # TODO: Manage the feature which use for managing exsist services.
                 logger.info(f"=====> Exsist Services Management <=====")
                 service_code = 1
                 for service in service_manager.services_list:
@@ -56,18 +69,25 @@ def main():
                 # logger.info("=====> a. Select All")
                 logger.info("=====> q. Quit")
 
-                manage_choice = input(f"Please select a chioce to carry out: ")
+                manage_choice = input("Please select a choice to carry out: ").strip()
                 if manage_choice == 'q':
                     logger.info("Exit main manager")
                     break
                 # if manage_choice == 'a':
                 #     pass
-                while int(manage_choice):
+                try:
                     code = int(manage_choice)
-                    service_manager.services_list[code].service_operation()
-                    break
+                    if 0 <= code < len(service_manager.services_list):
+                        service_manager.services_list[code].service_operation()
+                    else:
+                        logger.error("Invalid service code.")
+                except ValueError:
+                    logger.error("Please enter a valid number.")
         if start_choice == 'c':
             service_manager.list_services()
         if start_choice == 'q':
             logger.info(f"Exit the manager process")
             break
+
+if __name__ == '__main__':
+    sys.exit(main_menu())