import unittest
import argparse
import sys

# Define the test file names
test_files = [
    "testHistory.py",
    "testLayer.py",
    "testLRscheduler.py",
    "testModel.py",
    "testSimulationScheduler.py",
    "testSimulationScore.py",
    "testStorage.py",
    "unittests_autogenerated.py",
    "unittests_basic.py",
    "testTrain.py",
    "testOptimizers.py"
]

# Argument parser for CPU/GPU switch
parser = argparse.ArgumentParser(description="Run test suite with CPU/GPU mode.")
parser.add_argument('--mode', choices=['cpu', 'gpu'], default='cpu', help="Select computation mode: 'cpu' or 'gpu'")
args = parser.parse_args()
mode = args.mode

# Create a test loader
loader = unittest.TestLoader()

# Create an empty test suite
suite = unittest.TestSuite()

# Add tests from each test file to the test suite
for test_file in test_files:
    # Discover tests in the test file
    discovered_tests = loader.discover('./tests/', pattern=test_file)
    suite.addTests(discovered_tests)

# Create a test runner that will output the results to the console
runner = unittest.TextTestRunner()

# Run the test suite and exit with appropriate exit code
if __name__ == '__main__':
    result = runner.run(suite)
    if result.wasSuccessful():
        sys.exit(0)  # Exit with code 0 if all tests passed successfully
    else:
        sys.exit(1)  # Exit with code 1 if there were failures or errors in tests
