feat(performance-metrics): store to file

performance-metrics/src/performance_metrics/datashapes.py

@@ -2,6 +2,7 @@
 
 from enum import Enum
 import dataclasses
+from typing import Tuple
 
 
 class RobotContextState(Enum):
@@ -37,26 +38,30 @@ def from_id(cls, state_id: int) -> "RobotContextState":
 
 
 @dataclasses.dataclass(frozen=True)
-class RawDurationData:
-    """Represents raw duration data for a process or function.
+class RawContextData:
+    """Represents raw duration data with context state information.
 
     Attributes:
     - function_start_time (int): The start time of the function.
     - duration_measurement_start_time (int): The start time for duration measurement.
     - duration_measurement_end_time (int): The end time for duration measurement.
+    - state (RobotContextStates): The current state of the context.
     """
 
     func_start: int
     duration_start: int
     duration_end: int
-
-
-@dataclasses.dataclass(frozen=True)
-class RawContextData(RawDurationData):
-    """Extends RawDurationData with context state information.
-
-    Attributes:
-    - state (RobotContextStates): The current state of the context.
-    """
-
     state: RobotContextState
+
+    @classmethod
+    def headers(self) -> Tuple[str, str, str]:
+        """Returns the headers for the raw context data."""
+        return ("state_id", "function_start_time", "duration")
+
+    def csv_row(self) -> Tuple[int, int, int]:
+        """Returns the raw context data as a string."""
+        return (
+            self.state.state_id,
+            self.func_start,
+            self.duration_end - self.duration_start,
+        )

performance-metrics/src/performance_metrics/robot_context_tracker.py

@@ -1,5 +1,9 @@
 """Module for tracking robot context and execution duration for different operations."""
 
+import csv
+from pathlib import Path
+import os
+
 from functools import wraps
 from time import perf_counter_ns, clock_gettime_ns, CLOCK_REALTIME
 from typing import Callable, TypeVar
@@ -17,9 +21,10 @@
 class RobotContextTracker:
     """Tracks and stores robot context and execution duration for different operations."""
 
-    def __init__(self, should_track: bool = False) -> None:
+    def __init__(self, storage_file_path: Path, should_track: bool = False) -> None:
         """Initializes the RobotContextTracker with an empty storage list."""
         self._storage: deque[RawContextData] = deque()
+        self._storage_file_path = storage_file_path
         self._should_track = should_track
 
     def track(self, state: RobotContextState) -> Callable:  # type: ignore
@@ -57,3 +62,14 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:
             return wrapper
 
         return inner_decorator
+
+    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)

performance-metrics/tests/performance_metrics/test_robot_context_tracker.py

@@ -1,6 +1,7 @@
 """Tests for the RobotContextTracker class in performance_metrics.robot_context_tracker."""
 
 import asyncio
+from pathlib import Path
 import pytest
 from performance_metrics.robot_context_tracker import RobotContextTracker
 from performance_metrics.datashapes import RobotContextState
@@ -15,9 +16,9 @@
 
 
 @pytest.fixture
-def robot_context_tracker() -> RobotContextTracker:
+def robot_context_tracker(tmp_path: Path) -> RobotContextTracker:
     """Fixture to provide a fresh instance of RobotContextTracker for each test."""
-    return RobotContextTracker(should_track=True)
+    return RobotContextTracker(storage_file_path=tmp_path, should_track=True)
 
 
 def test_robot_context_tracker(robot_context_tracker: RobotContextTracker) -> None:
@@ -205,9 +206,9 @@ async def second_async_calibrating() -> None:
     ), "All tracked operations should be in CALIBRATING state."
 
 
-def test_no_tracking() -> None:
+def test_no_tracking(tmp_path: Path) -> None:
     """Tests that operations are not tracked when tracking is disabled."""
-    robot_context_tracker = RobotContextTracker(should_track=False)
+    robot_context_tracker = RobotContextTracker(tmp_path, should_track=False)
 
     @robot_context_tracker.track(state=RobotContextState.STARTING_UP)
     def operation_without_tracking() -> None:
@@ -218,3 +219,33 @@ def operation_without_tracking() -> None:
     assert (
         len(robot_context_tracker._storage) == 0
     ), "Operation should not be tracked when tracking is disabled."
+
+
+async def test_storing_to_file(tmp_path: Path) -> None:
+    """Tests storing the tracked data to a file."""
+    file_path = tmp_path / "test_file.csv"
+    robot_context_tracker = RobotContextTracker(file_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(file_path, "r") as file:
+        lines = file.readlines()
+        assert (
+            len(lines) == 4
+        ), "All stored data + header should be written to the file."