diff --git a/api/src/opentrons/cli/analyze.py b/api/src/opentrons/cli/analyze.py index b4f4f12ef321..42ca29a2b81e 100644 --- a/api/src/opentrons/cli/analyze.py +++ b/api/src/opentrons/cli/analyze.py @@ -28,9 +28,7 @@ ) from opentrons_shared_data.robot.dev_types import RobotType -from opentrons import get_robot_context_tracker - -_robot_context_tracker = get_robot_context_tracker() +from opentrons.util.performance_helpers import track_analysis @click.command() @@ -66,7 +64,7 @@ def _get_input_files(files_and_dirs: Sequence[Path]) -> List[Path]: return results -@_robot_context_tracker.track_analysis() +@track_analysis async def _analyze( files_and_dirs: Sequence[Path], json_output: Optional[AsyncPath], diff --git a/api/src/opentrons/util/performance_helpers.py b/api/src/opentrons/util/performance_helpers.py new file mode 100644 index 000000000000..12d7ea3f44f2 --- /dev/null +++ b/api/src/opentrons/util/performance_helpers.py @@ -0,0 +1,73 @@ +"""Performance helpers for tracking robot context.""" + +from pathlib import Path +from opentrons_shared_data.performance.dev_types import ( + SupportsTracking, + F, + RobotContextState, +) +from opentrons_shared_data.robot.dev_types import RobotTypeEnum +from opentrons.config import ( + get_performance_metrics_data_dir, + robot_configs, + feature_flags as ff, +) +from typing import Callable, Type + +performance_metrics_dir: Path = get_performance_metrics_data_dir() +should_track: bool = ff.enable_performance_metrics( + RobotTypeEnum.robot_literal_to_enum(robot_configs.load().model) +) + + +def _handle_package_import() -> Type[SupportsTracking]: + """Handle the import of the performance_metrics package. + + If the package is not available, return a stubbed tracker. + """ + try: + from performance_metrics import RobotContextTracker + + return RobotContextTracker + except ImportError: + return StubbedTracker + + +package_to_use = _handle_package_import() +_robot_context_tracker: SupportsTracking | None = None + + +class StubbedTracker(SupportsTracking): + """A stubbed tracker that does nothing.""" + + def __init__(self, storage_dir: Path, should_track: bool) -> None: + """Initialize the stubbed tracker.""" + pass + + def track(self, state: RobotContextState) -> Callable[[F], F]: + """Return the function unchanged.""" + + def inner_decorator(func: F) -> F: + """Return the function unchanged.""" + return func + + return inner_decorator + + def store(self) -> None: + """Do nothing.""" + pass + + +def _get_robot_context_tracker() -> SupportsTracking: + """Singleton for the robot context tracker.""" + global _robot_context_tracker + if _robot_context_tracker is None: + _robot_context_tracker = package_to_use(performance_metrics_dir, should_track) + return _robot_context_tracker + + +def track_analysis(func: F) -> F: + """Track the analysis of a protocol.""" + return _get_robot_context_tracker().track(RobotContextState.ANALYZING_PROTOCOL)( + func + ) diff --git a/api/tests/opentrons/cli/test_cli.py b/api/tests/opentrons/cli/test_cli.py index 7592fa940d07..007a7dd6a030 100644 --- a/api/tests/opentrons/cli/test_cli.py +++ b/api/tests/opentrons/cli/test_cli.py @@ -11,13 +11,15 @@ import pytest from click.testing import CliRunner -from opentrons import get_robot_context_tracker +from opentrons.util.performance_helpers import _get_robot_context_tracker # Enable tracking for the RobotContextTracker # This must come before the import of the analyze CLI -context_tracker = get_robot_context_tracker() -context_tracker._should_track = True +context_tracker = _get_robot_context_tracker() + +# Ignore the type error for the next line, as we're setting a private attribute for testing purposes +context_tracker._should_track = True # type: ignore[attr-defined] from opentrons.cli.analyze import analyze # noqa: E402 @@ -253,7 +255,7 @@ def test_python_error_line_numbers( assert error["detail"] == expected_detail -def test_tracking_of_analyis_with_robot_context_tracker(tmp_path: Path) -> None: +def test_track_analysis(tmp_path: Path) -> None: """Test that the RobotContextTracker tracks analysis.""" protocol_source = textwrap.dedent( """ @@ -267,8 +269,8 @@ def run(protocol): protocol_source_file = tmp_path / "protocol.py" protocol_source_file.write_text(protocol_source, encoding="utf-8") - before_analysis = len(context_tracker._storage) + before_analysis = len(context_tracker._storage) # type: ignore[attr-defined] _get_analysis_result([protocol_source_file]) - assert len(context_tracker._storage) == before_analysis + 1 + assert len(context_tracker._storage) == before_analysis + 1 # type: ignore[attr-defined] diff --git a/api/tests/opentrons/util/test_performance_helpers.py b/api/tests/opentrons/util/test_performance_helpers.py new file mode 100644 index 000000000000..57a42ef6a714 --- /dev/null +++ b/api/tests/opentrons/util/test_performance_helpers.py @@ -0,0 +1,28 @@ +"""Tests for performance_helpers.""" + +from pathlib import Path +from opentrons_shared_data.performance.dev_types import RobotContextState +from opentrons.util.performance_helpers import ( + StubbedTracker, + _get_robot_context_tracker, +) + + +def test_return_function_unchanged() -> None: + """Test that the function is returned unchanged when using StubbedTracker.""" + tracker = StubbedTracker(Path("/path/to/storage"), True) + + def func_to_track() -> None: + pass + + assert ( + tracker.track(RobotContextState.ANALYZING_PROTOCOL)(func_to_track) + is func_to_track + ) + + +def test_singleton_tracker() -> None: + """Test that the tracker is a singleton.""" + tracker = _get_robot_context_tracker() + tracker2 = _get_robot_context_tracker() + assert tracker is tracker2