#!/usr/bin/env python3 # # Copyright 2025 The InfiniFlow Authors. All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import sys import os import argparse import subprocess from pathlib import Path from typing import List class Colors: """ANSI color codes for terminal output""" RED = '\033[0;31m' GREEN = '\033[0;32m' YELLOW = '\033[1;33m' BLUE = '\033[0;34m' NC = '\033[0m' # No Color class TestRunner: """RAGFlow Unit Test Runner""" def __init__(self): self.project_root = Path(__file__).parent.resolve() self.ut_dir = Path(self.project_root / 'test' / 'unit_test') # Default options self.coverage = False self.parallel = False self.verbose = False self.markers = "" # Python interpreter path self.python = sys.executable @staticmethod def print_info(message: str) -> None: """Print informational message""" print(f"{Colors.BLUE}[INFO]{Colors.NC} {message}") @staticmethod def print_error(message: str) -> None: """Print error message""" print(f"{Colors.RED}[ERROR]{Colors.NC} {message}") @staticmethod def show_usage() -> None: """Display usage information""" usage = """ RAGFlow Unit Test Runner Usage: python run_tests.py [OPTIONS] OPTIONS: -h, --help Show this help message -c, --coverage Run tests with coverage report -p, --parallel Run tests in parallel (requires pytest-xdist) -v, --verbose Verbose output -t, --test FILE Run specific test file or directory -m, --markers MARKERS Run tests with specific markers (e.g., "unit", "integration") EXAMPLES: # Run all tests python run_tests.py # Run with coverage python run_tests.py --coverage # Run in parallel python run_tests.py --parallel # Run specific test file python run_tests.py --test services/test_dialog_service.py # Run only unit tests python run_tests.py --markers "unit" # Run tests with coverage and parallel execution python run_tests.py --coverage --parallel """ print(usage) def build_pytest_command(self) -> List[str]: """Build the pytest command arguments""" cmd = ["pytest", str(self.ut_dir)] # Add test path # Add markers if self.markers: cmd.extend(["-m", self.markers]) # Add verbose flag if self.verbose: cmd.extend(["-vv"]) else: cmd.append("-v") # Add coverage if self.coverage: # Relative path from test directory to source code source_path = str(self.project_root / "common") cmd.extend([ "--cov", source_path, "--cov-report", "html", "--cov-report", "term" ]) # Add parallel execution if self.parallel: # Try to get number of CPU cores try: import multiprocessing cpu_count = multiprocessing.cpu_count() cmd.extend(["-n", str(cpu_count)]) except ImportError: # Fallback to auto if multiprocessing not available cmd.extend(["-n", "auto"]) # Add default options from pyproject.toml if it exists pyproject_path = self.project_root / "pyproject.toml" if pyproject_path.exists(): cmd.extend(["--config-file", str(pyproject_path)]) return cmd def run_tests(self) -> bool: """Execute the pytest command""" # Change to test directory os.chdir(self.project_root) # Build command cmd = self.build_pytest_command() # Print test configuration self.print_info("Running RAGFlow Unit Tests") self.print_info("=" * 40) self.print_info(f"Test Directory: {self.ut_dir}") self.print_info(f"Coverage: {self.coverage}") self.print_info(f"Parallel: {self.parallel}") self.print_info(f"Verbose: {self.verbose}") if self.markers: self.print_info(f"Markers: {self.markers}") print(f"\n{Colors.BLUE}[EXECUTING]{Colors.NC} {' '.join(cmd)}\n") # Run pytest try: result = subprocess.run(cmd, check=False) if result.returncode == 0: print(f"\n{Colors.GREEN}[SUCCESS]{Colors.NC} All tests passed!") if self.coverage: coverage_dir = self.ut_dir / "htmlcov" if coverage_dir.exists(): index_file = coverage_dir / "index.html" print(f"\n{Colors.BLUE}[INFO]{Colors.NC} Coverage report generated:") print(f" {index_file}") print("\nOpen with:") print(f" - Windows: start {index_file}") print(f" - macOS: open {index_file}") print(f" - Linux: xdg-open {index_file}") return True else: print(f"\n{Colors.RED}[FAILURE]{Colors.NC} Some tests failed!") return False except KeyboardInterrupt: print(f"\n{Colors.YELLOW}[INTERRUPTED]{Colors.NC} Test execution interrupted by user") return False except Exception as e: self.print_error(f"Failed to execute tests: {e}") return False def parse_arguments(self) -> bool: """Parse command line arguments""" parser = argparse.ArgumentParser( description="RAGFlow Unit Test Runner", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: python run_tests.py # Run all tests python run_tests.py --coverage # Run with coverage python run_tests.py --parallel # Run in parallel python run_tests.py --test services/test_dialog_service.py # Run specific test python run_tests.py --markers "unit" # Run only unit tests """ ) parser.add_argument( "-c", "--coverage", action="store_true", help="Run tests with coverage report" ) parser.add_argument( "-p", "--parallel", action="store_true", help="Run tests in parallel (requires pytest-xdist)" ) parser.add_argument( "-v", "--verbose", action="store_true", help="Verbose output" ) parser.add_argument( "-t", "--test", type=str, default="", help="Run specific test file or directory" ) parser.add_argument( "-m", "--markers", type=str, default="", help="Run tests with specific markers (e.g., 'unit', 'integration')" ) try: args = parser.parse_args() # Set options self.coverage = args.coverage self.parallel = args.parallel self.verbose = args.verbose self.markers = args.markers return True except SystemExit: # argparse already printed help, just exit return False except Exception as e: self.print_error(f"Error parsing arguments: {e}") return False def run(self) -> int: """Main execution method""" # Parse command line arguments if not self.parse_arguments(): return 1 # Run tests success = self.run_tests() return 0 if success else 1 def main(): """Entry point""" runner = TestRunner() return runner.run() if __name__ == "__main__": sys.exit(main())