Add long event tracker

samcli/lib/telemetry/event.py

@@ -173,5 +173,79 @@ def send_events():
             EventTracker._events = []  # Manual clear_trackers() since we're within the lock
 
 
+def track_long_event(start_event_name: str, start_event_value: str, end_event_name: str, end_event_value: str):
+    """Decorator for tracking events that occur at start and end of a function.
+
+    The decorator tracks two Events total, where the first Event occurs
+    at the start of the decorated function's execution (prior to its
+    first line) and the second Event occurs after the function has ended
+    (after the final line of the function has executed).
+
+    Parameters
+    ----------
+    start_event_name: str
+        The name of the Event that is executed at the start of the
+        decorated function's execution. Must be a valid EventName
+        value or the decorator will not run.
+    start_event_value: str
+        The value of the Event that is executed at the start of the
+        decorated function's execution. Must be a valid EventType
+        value for the passed `start_event_name` or the decorator
+        will not run.
+    end_event_name: str
+        The name of the Event that is executed at the end of the
+        decorated function's execution. Must be a valid EventName
+        value or the decorator will not run.
+    end_event_value: str
+        The value of the Event that is executed at the end of the
+        decorated function's execution. Must be a valid EventType
+        value for the passed `end_event_name` or the decorator
+        will not run.
+
+    Examples
+    --------
+    >>> @track_long_event("FuncStart", "Func1", "FuncEnd", "Func1")
+        def func1(...):
+            # do things
+
+    >>> @track_long_event("FuncStart", "Func2", "FuncEnd", "Func2")
+        def func2(...):
+            # do things
+    """
+    should_track = True
+    try:
+        # Check that passed values are valid Events
+        Event(start_event_name, start_event_value)
+        Event(end_event_name, end_event_value)
+    except EventCreationError as e:
+        LOG.debug("Error occurred while trying to track an event: %s\nDecorator not run.", e)
+        should_track = False
+
+    def decorator_for_events(func):
+        """The actual decorator"""
+
+        def wrapped(*args, **kwargs):
+            if should_track:
+                EventTracker.track_event(start_event_name, start_event_value)
+            exception = None
+
+            try:
+                return_value = func(*args, **kwargs)
+            except Exception as e:
+                exception = e
+
+            if should_track:
+                EventTracker.track_event(end_event_name, end_event_value)
+                EventTracker.send_events()  # Ensure Events are sent at the end of execution
+            if exception:
+                raise exception
+
+            return return_value
+
+        return wrapped
+
+    return decorator_for_events
+
+
 class EventCreationError(Exception):
     """Exception called when an Event is not properly created."""

tests/unit/lib/telemetry/test_event.py

@@ -4,10 +4,11 @@
 
 from enum import Enum
 import threading
+from typing import List, Tuple
 from unittest import TestCase
 from unittest.mock import ANY, Mock, patch
 
-from samcli.lib.telemetry.event import Event, EventCreationError, EventTracker
+from samcli.lib.telemetry.event import Event, EventCreationError, EventTracker, track_long_event
 
 
 class DummyEventName(Enum):
@@ -170,3 +171,63 @@ def make_mock_event(name, value):
         EventTracker.track_event("TheStrawThat", "BreaksTheCamel'sBack")
 
         send_mock.assert_called()
+
+
+class TestTrackLongEvent(TestCase):
+    @patch("samcli.lib.telemetry.event.EventTracker.send_events")
+    @patch("samcli.lib.telemetry.event.EventTracker.track_event")
+    @patch("samcli.lib.telemetry.event.Event", return_value=None)
+    def test_long_event_is_tracked(self, event_mock, track_mock, send_mock):
+        mock_tracker = {}
+        mock_tracker["tracked_events"]: List[Tuple[str, str]] = []  # Tuple to bypass Event verification
+        mock_tracker["emitted_events"]: List[Tuple[str, str]] = []
+
+        def mock_track(name, value):
+            mock_tracker["tracked_events"].append((name, value))
+
+        def mock_send():
+            mock_tracker["emitted_events"] = mock_tracker["tracked_events"]
+            mock_tracker["tracked_events"] = []  # Mimic clear_trackers()
+
+        track_mock.side_effect = mock_track
+        send_mock.side_effect = mock_send
+
+        @track_long_event("StartEvent", "StartValue", "EndEvent", "EndValue")
+        def func():
+            self.assertEqual(len(mock_tracker["tracked_events"]), 1, "Starting event not tracked.")
+            self.assertIn(("StartEvent", "StartValue"), mock_tracker["tracked_events"], "Incorrect starting event.")
+
+        func()
+
+        self.assertEqual(len(mock_tracker["tracked_events"]), 0, "Tracked events not reset; send_events not called.")
+        self.assertEqual(len(mock_tracker["emitted_events"]), 2, "Unexpected number of emitted events.")
+        self.assertIn(("StartEvent", "StartValue"), mock_tracker["emitted_events"], "Starting event not tracked.")
+        self.assertIn(("EndEvent", "EndValue"), mock_tracker["emitted_events"], "Ending event not tracked.")
+
+    @patch("samcli.lib.telemetry.event.EventTracker.send_events")
+    @patch("samcli.lib.telemetry.event.EventTracker.track_event")
+    def test_nothing_tracked_if_invalid_events(self, track_mock, send_mock):
+        mock_tracker = {}
+        mock_tracker["tracked_events"]: List[Tuple[str, str]] = []  # Tuple to bypass Event verification
+        mock_tracker["emitted_events"]: List[Tuple[str, str]] = []
+
+        def mock_track(name, value):
+            mock_tracker["tracked_events"].append((name, value))
+
+        def mock_send():
+            mock_tracker["emitted_events"] = mock_tracker["tracked_events"]
+            mock_tracker["tracked_events"] = []  # Mimic clear_trackers()
+
+        track_mock.side_effect = mock_track
+        send_mock.side_effect = mock_send
+
+        @track_long_event("DefinitelyNotARealEvent", "Nope", "ThisEventDoesntExist", "NuhUh")
+        def func():
+            self.assertEqual(len(mock_tracker["tracked_events"]), 0, "Events should not have been tracked.")
+
+        func()
+
+        self.assertEqual(len(mock_tracker["tracked_events"]), 0, "Events should not have been tracked.")
+        self.assertEqual(len(mock_tracker["emitted_events"]), 0, "Events should not have been emitted.")
+        track_mock.assert_not_called()  # Tracker should not have been called
+        send_mock.assert_not_called()  # Sender should not have been called