refactor(performance-metrics): separate storage logic from tracker

api/tests/opentrons/cli/test_cli.py

@@ -5,12 +5,15 @@
 import tempfile
 import textwrap
 
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Any, Dict, Iterator, List, Optional
 from pathlib import Path
 
 import pytest
 from click.testing import CliRunner
+from opentrons_shared_data.performance.dev_types import (
+    RobotContextState,
+)
 from opentrons.util.performance_helpers import _get_robot_context_tracker
 
 
@@ -24,6 +27,18 @@
 from opentrons.cli.analyze import analyze  # noqa: E402
 
 
+@pytest.fixture
+def override_data_store(tmp_path: Path) -> Iterator[None]:
+    """Override the data store metadata for the RobotContextTracker."""
+    old_store = context_tracker._store  # type: ignore[attr-defined]
+    old_metadata = old_store.metadata
+    new_metadata = replace(old_metadata, storage_dir=tmp_path)
+    context_tracker._store = old_store.__class__(metadata=new_metadata)  # type: ignore[attr-defined]
+    context_tracker._store.setup()  # type: ignore[attr-defined]
+    yield
+    context_tracker._store = old_store  # type: ignore[attr-defined]
+
+
 def _list_fixtures(version: int) -> Iterator[Path]:
     return Path(__file__).parent.glob(
         f"../../../../shared-data/protocol/fixtures/{version}/*.json"
@@ -255,6 +270,7 @@ def test_python_error_line_numbers(
     assert error["detail"] == expected_detail
 
 
+@pytest.mark.usefixtures("override_data_store")
 def test_track_analysis(tmp_path: Path) -> None:
     """Test that the RobotContextTracker tracks analysis."""
     protocol_source = textwrap.dedent(
@@ -265,12 +281,27 @@ def run(protocol):
             pass
         """
     )
-
     protocol_source_file = tmp_path / "protocol.py"
     protocol_source_file.write_text(protocol_source, encoding="utf-8")
+    store = context_tracker._store  # type: ignore[attr-defined]
 
-    before_analysis = len(context_tracker._storage)  # type: ignore[attr-defined]
+    num_storage_entities_before_analysis = len(store._data)
 
     _get_analysis_result([protocol_source_file])
 
-    assert len(context_tracker._storage) == before_analysis + 1  # type: ignore[attr-defined]
+    assert len(store._data) == num_storage_entities_before_analysis + 1
+
+    with open(store.metadata.data_file_location, "r") as f:
+        stored_data = f.readlines()
+        assert len(stored_data) == 0
+
+    context_tracker.store()
+
+    with open(store.metadata.data_file_location, "r") as f:
+        stored_data = f.readlines()
+        stored_data = [line.strip() for line in stored_data if line.strip()]
+        assert len(stored_data) == 1
+        state_id, start_time, duration = stored_data[0].strip().split(",")
+        assert state_id == str(RobotContextState.ANALYZING_PROTOCOL.state_id)
+        assert start_time.isdigit()
+        assert duration.isdigit()

performance-metrics/src/performance_metrics/datashapes.py

@@ -1,12 +1,32 @@
-"""Defines data classes and enums used in the performance metrics module."""
+"""Defines the shape of stored data."""
 
 import dataclasses
-from typing import Tuple
+from typing import Sequence, Tuple, Protocol, Union
 from opentrons_shared_data.performance.dev_types import RobotContextState
 
+StorableData = Union[int, float, str]
+
+
+class SupportsCSVStorage(Protocol):
+    """A protocol for classes that support CSV storage."""
+
+    @classmethod
+    def headers(self) -> Tuple[str, ...]:
+        """Returns the headers for the CSV data."""
+        ...
+
+    def csv_row(self) -> Tuple[StorableData, ...]:
+        """Returns the object as a CSV row."""
+        ...
+
+    @classmethod
+    def from_csv_row(cls, row: Tuple[StorableData, ...]) -> "SupportsCSVStorage":
+        """Returns an object from a CSV row."""
+        ...
+
 
 @dataclasses.dataclass(frozen=True)
-class RawContextData:
+class RawContextData(SupportsCSVStorage):
     """Represents raw duration data with context state information.
 
     Attributes:
@@ -16,10 +36,9 @@ class RawContextData:
     - state (RobotContextStates): The current state of the context.
     """
 
-    func_start: int
-    duration_start: int
-    duration_end: int
     state: RobotContextState
+    func_start: int
+    duration: int
 
     @classmethod
     def headers(self) -> Tuple[str, str, str]:
@@ -31,5 +50,14 @@ def csv_row(self) -> Tuple[int, int, int]:
         return (
             self.state.state_id,
             self.func_start,
-            self.duration_end - self.duration_start,
+            self.duration,
+        )
+
+    @classmethod
+    def from_csv_row(cls, row: Sequence[StorableData]) -> SupportsCSVStorage:
+        """Returns a RawContextData object from a CSV row."""
+        return cls(
+            state=RobotContextState.from_id(int(row[0])),
+            func_start=int(row[1]),
+            duration=int(row[2]),
         )

performance-metrics/src/performance_metrics/metrics_store.py

@@ -0,0 +1,37 @@
+"""Interface for storing performance metrics data to a CSV file."""
+
+import csv
+import typing
+from opentrons_shared_data.performance.dev_types import MetricsMetadata
+from performance_metrics.datashapes import SupportsCSVStorage
+
+T = typing.TypeVar("T", bound=SupportsCSVStorage)
+
+
+class MetricsStore(typing.Generic[T]):
+    """Dataclass to store data for tracking robot context."""
+
+    def __init__(self, metadata: MetricsMetadata) -> None:
+        """Initialize the metrics store."""
+        self.metadata = metadata
+        self._data: typing.List[T] = []
+
+    def add(self, context_data: T) -> None:
+        """Add data to the store."""
+        self._data.append(context_data)
+
+    def setup(self) -> None:
+        """Set up the data store."""
+        self.metadata.storage_dir.mkdir(parents=True, exist_ok=True)
+        self.metadata.data_file_location.touch(exist_ok=True)
+        self.metadata.headers_file_location.touch(exist_ok=True)
+        self.metadata.headers_file_location.write_text(",".join(self.metadata.headers))
+
+    def store(self) -> None:
+        """Clear the stored data and write it to the storage file."""
+        stored_data = self._data.copy()
+        self._data.clear()
+        rows_to_write = [context_data.csv_row() for context_data in stored_data]
+        with open(self.metadata.data_file_location, "a") as storage_file:
+            writer = csv.writer(storage_file)
+            writer.writerows(rows_to_write)

performance-metrics/src/performance_metrics/robot_context_tracker.py

@@ -1,23 +1,22 @@
 """Module for tracking robot context and execution duration for different operations."""
 
-import csv
 from pathlib import Path
 import platform
 
 from functools import wraps, partial
 from time import perf_counter_ns
-import os
-from typing import Callable, TypeVar, cast
+from typing import Callable, TypeVar, cast, Literal, Final
 
 
 from typing_extensions import ParamSpec
-from collections import deque
 from performance_metrics.datashapes import (
     RawContextData,
 )
+from performance_metrics.metrics_store import MetricsStore
 from opentrons_shared_data.performance.dev_types import (
     RobotContextState,
     SupportsTracking,
+    MetricsMetadata,
 )
 
 P = ParamSpec("P")
@@ -47,14 +46,22 @@ def _get_timing_function() -> Callable[[], int]:
 class RobotContextTracker(SupportsTracking):
     """Tracks and stores robot context and execution duration for different operations."""
 
-    FILE_NAME = "context_data.csv"
+    METADATA_NAME: Final[Literal["robot_context_data"]] = "robot_context_data"
 
     def __init__(self, storage_location: Path, should_track: bool = False) -> None:
         """Initializes the RobotContextTracker with an empty storage list."""
-        self._storage: deque[RawContextData] = deque()
-        self._storage_file_path = storage_location / self.FILE_NAME
+        self._store = MetricsStore[RawContextData](
+            MetricsMetadata(
+                name=self.METADATA_NAME,
+                storage_dir=storage_location,
+                headers=RawContextData.headers(),
+            )
+        )
         self._should_track = should_track
 
+        if self._should_track:
+            self._store.setup()
+
     def track(self, state: RobotContextState) -> Callable:  # type: ignore
         """Decorator factory for tracking the execution duration and state of robot operations.
 
@@ -77,14 +84,14 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
                     result = func(*args, **kwargs)
                 finally:
                     duration_end_time = perf_counter_ns()
-                    self._storage.append(
+                    self._store.add(
                         RawContextData(
-                            function_start_time,
-                            duration_start_time,
-                            duration_end_time,
-                            state,
+                            func_start=function_start_time,
+                            duration=duration_end_time - duration_start_time,
+                            state=state,
                         )
                     )
+
                 return result
 
             return wrapper
@@ -93,11 +100,6 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
 
     def store(self) -> None:
         """Returns the stored context data and clears the storage list."""
-        stored_data = self._storage.copy()
-        self._storage.clear()
-        rows_to_write = [context_data.csv_row() for context_data in stored_data]
-        os.makedirs(self._storage_file_path.parent, exist_ok=True)
-        with open(self._storage_file_path, "a") as storage_file:
-            writer = csv.writer(storage_file)
-            writer.writerow(RawContextData.headers())
-            writer.writerows(rows_to_write)
+        if not self._should_track:
+            return
+        self._store.store()

performance-metrics/tests/performance_metrics/test_metrics_store.py

@@ -0,0 +1,52 @@
+"""Tests for the metrics store."""
+
+from pathlib import Path
+from time import sleep
+
+from opentrons_shared_data.performance.dev_types import RobotContextState
+from performance_metrics.datashapes import RawContextData
+from performance_metrics.robot_context_tracker import RobotContextTracker
+
+# Corrected times in seconds
+STARTING_TIME = 0.001
+CALIBRATING_TIME = 0.002
+ANALYZING_TIME = 0.003
+RUNNING_TIME = 0.004
+SHUTTING_DOWN_TIME = 0.005
+
+
+async def test_storing_to_file(tmp_path: Path) -> None:
+    """Tests storing the tracked data to a file."""
+    robot_context_tracker = RobotContextTracker(tmp_path, should_track=True)
+
+    @robot_context_tracker.track(state=RobotContextState.STARTING_UP)
+    def starting_robot() -> None:
+        sleep(STARTING_TIME)
+
+    @robot_context_tracker.track(state=RobotContextState.CALIBRATING)
+    def calibrating_robot() -> None:
+        sleep(CALIBRATING_TIME)
+
+    @robot_context_tracker.track(state=RobotContextState.ANALYZING_PROTOCOL)
+    def analyzing_protocol() -> None:
+        sleep(ANALYZING_TIME)
+
+    starting_robot()
+    calibrating_robot()
+    analyzing_protocol()
+
+    robot_context_tracker.store()
+
+    with open(robot_context_tracker._store.metadata.data_file_location, "r") as file:
+        lines = file.readlines()
+        assert len(lines) == 3, "All stored data should be written to the file."
+
+        split_lines: list[list[str]] = [line.strip().split(",") for line in lines]
+        assert all(
+            RawContextData.from_csv_row(line) for line in split_lines
+        ), "All lines should be valid RawContextData instances."
+
+    with open(robot_context_tracker._store.metadata.headers_file_location, "r") as file:
+        headers = file.readlines()
+        assert len(headers) == 1, "Header should be written to the headers file."
+        assert tuple(headers[0].strip().split(",")) == RawContextData.headers()