feat(consume): track tests by pre-alloc group for client cleanup

Author: danceratopz

src/pytest_plugins/consume/simulators/enginex/conftest.py

@@ -4,19 +4,24 @@
 Configures the hive back-end & EL clients for test execution with BlockchainEngineXFixtures.
 """
 
+import logging
+
 import pytest
 from hive.client import Client
 
 from ethereum_test_exceptions import ExceptionMapper
 from ethereum_test_fixtures import BlockchainEngineXFixture
 from ethereum_test_rpc import EngineRPC
 
+logger = logging.getLogger(__name__)
+
 pytest_plugins = (
     "pytest_plugins.consume.simulators.base",
     "pytest_plugins.consume.simulators.multi_test_client",
     "pytest_plugins.consume.simulators.test_case_description",
     "pytest_plugins.consume.simulators.timing_data",
     "pytest_plugins.consume.simulators.exceptions",
+    "pytest_plugins.consume.simulators.helpers.test_tracker",
 )
 
 
@@ -40,6 +45,54 @@ def test_suite_description() -> str:
     )
 
 
+def pytest_collection_modifyitems(session, config, items):
+    """
+    Build pre-allocation group test counts during collection phase.
+
+    This hook analyzes all collected test items to determine how many tests
+    belong to each pre-allocation group, enabling automatic client cleanup
+    when all tests in a group are complete.
+    """
+    # Only process items for enginex simulator
+    if not hasattr(config, "_supported_fixture_formats"):
+        return
+
+    if BlockchainEngineXFixture.format_name not in config._supported_fixture_formats:
+        return
+
+    # Get the test tracker from session if available
+    test_tracker = getattr(session, "_pre_alloc_group_test_tracker", None)
+    if test_tracker is None:
+        # Tracker will be created later by the fixture, store counts on session for now
+        group_counts = {}
+        for item in items:
+            if hasattr(item, "callspec") and "test_case" in item.callspec.params:
+                test_case = item.callspec.params["test_case"]
+                if hasattr(test_case, "pre_hash"):
+                    pre_hash = test_case.pre_hash
+                    group_counts[pre_hash] = group_counts.get(pre_hash, 0) + 1
+
+        # Store on session for later retrieval by test_tracker fixture
+        session._pre_alloc_group_counts = group_counts
+        logger.info(
+            f"Collected {len(group_counts)} pre-allocation groups with tests: {dict(group_counts)}"
+        )
+    else:
+        # Update tracker directly if it exists
+        group_counts = {}
+        for item in items:
+            if hasattr(item, "callspec") and "test_case" in item.callspec.params:
+                test_case = item.callspec.params["test_case"]
+                if hasattr(test_case, "pre_hash"):
+                    pre_hash = test_case.pre_hash
+                    group_counts[pre_hash] = group_counts.get(pre_hash, 0) + 1
+
+        for pre_hash, count in group_counts.items():
+            test_tracker.set_group_test_count(pre_hash, count)
+
+        logger.info(f"Updated test tracker with {len(group_counts)} pre-allocation groups")
+
+
 @pytest.fixture(scope="function")
 def engine_rpc(client: Client, client_exception_mapper: ExceptionMapper | None) -> EngineRPC:
     """Initialize engine RPC client for the execution client under test."""

src/pytest_plugins/consume/simulators/helpers/client_wrapper.py

@@ -5,7 +5,7 @@
 import logging
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Dict, Optional, cast
+from typing import TYPE_CHECKING, Dict, Optional, cast
 
 from hive.client import Client, ClientType
 
@@ -15,6 +15,9 @@
 from ethereum_test_forks import Fork
 from pytest_plugins.consume.simulators.helpers.ruleset import ruleset
 
+if TYPE_CHECKING:
+    from .test_tracker import PreAllocGroupTestTracker
+
 logger = logging.getLogger(__name__)
 
 
@@ -248,12 +251,22 @@ def set_client(self, client: Client) -> None:
         )
 
     def stop(self) -> None:
-        """Override to log with pre_hash information."""
+        """Override to log with pre_hash information and actually stop the client."""
         if self._is_started:
             logger.info(
-                f"Marking multi-test client ({self.client_type.name}) for pre-allocation group "
-                f"{self.pre_hash} as stopped after {self.test_count} tests"
+                f"Stopping multi-test client ({self.client_type.name}) for pre-allocation group "
+                f"{self.pre_hash} after {self.test_count} tests"
             )
+            # Actually stop the Hive client
+            if self.client is not None:
+                try:
+                    self.client.stop()
+                    logger.debug(f"Hive client stopped for pre-allocation group {self.pre_hash}")
+                except Exception as e:
+                    logger.error(
+                        f"Error stopping Hive client for pre-allocation group {self.pre_hash}: {e}"
+                    )
+
             self.client = None
             self._is_started = False
 
@@ -283,6 +296,7 @@ def __init__(self) -> None:
 
         self.multi_test_clients: Dict[str, MultiTestClient] = {}
         self.pre_alloc_path: Optional[Path] = None
+        self.test_tracker: Optional["PreAllocGroupTestTracker"] = None
         self._initialized = True
         logger.info("MultiTestClientManager initialized")
 
@@ -297,6 +311,17 @@ def set_pre_alloc_path(self, path: Path) -> None:
         self.pre_alloc_path = path
         logger.debug(f"Pre-alloc path set to: {path}")
 
+    def set_test_tracker(self, test_tracker: "PreAllocGroupTestTracker") -> None:
+        """
+        Set the test tracker for automatic client cleanup.
+
+        Args:
+            test_tracker: The test tracker instance
+
+        """
+        self.test_tracker = test_tracker
+        logger.debug("Test tracker set for automatic client cleanup")
+
     def load_pre_alloc_group(self, pre_hash: str) -> PreAllocGroup:
         """
         Load the pre-allocation group for a given preHash.
@@ -373,12 +398,15 @@ def get_or_create_multi_test_client(
 
         return multi_test_client
 
-    def get_client_for_test(self, pre_hash: str) -> Optional[Client]:
+    def get_client_for_test(
+        self, pre_hash: str, test_id: Optional[str] = None
+    ) -> Optional[Client]:
         """
         Get the actual client instance for a test with the given preHash.
 
         Args:
             pre_hash: The hash identifying the pre-allocation group
+            test_id: Optional test ID for completion tracking
 
         Returns:
             The client instance if available, None otherwise
@@ -391,6 +419,58 @@ def get_client_for_test(self, pre_hash: str) -> Optional[Client]:
                 return multi_test_client.client
         return None
 
+    def mark_test_completed(self, pre_hash: str, test_id: str) -> None:
+        """
+        Mark a test as completed and trigger automatic client cleanup if appropriate.
+
+        Args:
+            pre_hash: The hash identifying the pre-allocation group
+            test_id: The unique test identifier
+
+        """
+        if self.test_tracker is None:
+            logger.debug("No test tracker available, skipping completion tracking")
+            return
+
+        # Mark test as completed in tracker
+        is_group_complete = self.test_tracker.mark_test_completed(pre_hash, test_id)
+
+        if is_group_complete:
+            # All tests in this pre-allocation group are complete
+            self._auto_stop_client_if_complete(pre_hash)
+
+    def _auto_stop_client_if_complete(self, pre_hash: str) -> None:
+        """
+        Automatically stop the client for a pre-allocation group if all tests are complete.
+
+        Args:
+            pre_hash: The hash identifying the pre-allocation group
+
+        """
+        if pre_hash not in self.multi_test_clients:
+            logger.debug(f"No client found for pre-allocation group {pre_hash}")
+            return
+
+        multi_test_client = self.multi_test_clients[pre_hash]
+        if not multi_test_client.is_running:
+            logger.debug(f"Client for pre-allocation group {pre_hash} is already stopped")
+            return
+
+        # Stop the client and remove from tracking
+        logger.info(
+            f"Auto-stopping client for pre-allocation group {pre_hash} - "
+            f"all tests completed ({multi_test_client.test_count} tests executed)"
+        )
+
+        try:
+            multi_test_client.stop()
+        except Exception as e:
+            logger.error(f"Error auto-stopping client for pre-allocation group {pre_hash}: {e}")
+        finally:
+            # Remove from tracking to free memory
+            del self.multi_test_clients[pre_hash]
+            logger.debug(f"Removed completed client from tracking: {pre_hash}")
+
     def stop_all_clients(self) -> None:
         """Mark all multi-test clients as stopped."""
         logger.info(f"Marking all {len(self.multi_test_clients)} multi-test clients as stopped")
@@ -422,4 +502,5 @@ def reset(self) -> None:
         self.stop_all_clients()
         self.multi_test_clients.clear()
         self.pre_alloc_path = None
+        self.test_tracker = None
         logger.info("MultiTestClientManager reset")