doc: add type hints for test schedular

Author: darwintree

tests/test_utils/test_scheduler.py

@@ -5,12 +5,13 @@
 import os
 import tomllib
 import time
+from typing import Callable
 
 SPECIAL_SCRIPTS = {
     "ghast_consensus_test.py": 1
 }
 
-def get_num_test_nodes(py, test_dir, script):
+def get_num_test_nodes(py: str, test_dir: str, script: str) -> int:
     if script in SPECIAL_SCRIPTS:
         return SPECIAL_SCRIPTS[script]
     toml_output = subprocess.check_output(
@@ -22,7 +23,15 @@ def get_num_test_nodes(py, test_dir, script):
 class TestScheduler:
     """Scheduler for managing test execution and controlling concurrency based on resource requirements"""
 
-    def __init__(self, task_executable, py, test_dir, max_workers, available_nodes, port_min, port_max, conflux_binary):
+    def __init__(self,
+                 task_executable: Callable[[str, str, str, int, int, int, str], None],
+                 py: str,
+                 test_dir: str,
+                 max_workers: int,
+                 available_nodes: int,
+                 port_min: int,
+                 port_max: int,
+                 conflux_binary: str):
         self.task_executable = task_executable
         self.py = py
         self.test_dir = test_dir
@@ -41,7 +50,7 @@ def __init__(self, task_executable, py, test_dir, max_workers, available_nodes,
         self.results = []
         self.failed_tests = set()
 
-    def schedule(self, test_scripts):
+    def schedule(self, test_scripts: list[str]) -> set[str]:
         """Schedules the execution of test scripts"""
 
         # Prepare task queue
@@ -53,7 +62,7 @@ def schedule(self, test_scripts):
             self._collect_results()
         return self.failed_tests
 
-    def _prepare_task_queue(self, test_scripts):
+    def _prepare_task_queue(self, test_scripts: list[str]) -> queue.Queue[tuple[str, int, int]]:
         """Prepares a task queue with scripts and resource requirements"""
 
         task_queue = queue.Queue()
@@ -77,7 +86,7 @@ def _prepare_task_queue(self, test_scripts):
         print(" Done")
         return task_queue
 
-    def _process_task_queue(self, executor, task_queue):
+    def _process_task_queue(self, executor: ThreadPoolExecutor, task_queue: queue.Queue[tuple[str, int, int]]):
         """Processes the task queue, scheduling tests based on resource availability"""
 
         while not task_queue.empty():
@@ -106,7 +115,7 @@ def _process_task_queue(self, executor, task_queue):
             except queue.Empty:
                 break
 
-    def _try_acquire_resources(self, nodes_needed):
+    def _try_acquire_resources(self, nodes_needed: int) -> bool:
         """Attempts to acquire required resources, returns True if successful"""
 
         with self.resource_lock:
@@ -116,15 +125,15 @@ def _try_acquire_resources(self, nodes_needed):
                 return True
             return False
 
-    def _release_resources(self, nodes_count):
+    def _release_resources(self, nodes_count: int):
         """Releases resources and notifies waiting threads"""
 
         with self.resource_lock:
             self.available_nodes += nodes_count
             self.available_workers += 1
             self.resource_event.set()
 
-    def _run_test_with_cleanup(self, script, index, nodes_count):
+    def _run_test_with_cleanup(self, script: str, index: int, nodes_count: int):
         """Runs a test and ensures resources are released"""
 
         try: