From 341202633115c746a70d4f44fb2018ecebb0aa22 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20S=C3=A1nchez-Gallego?= <gallegoj@uw.edu>
Date: Sat, 21 Dec 2024 23:38:24 -0800
Subject: [PATCH] Added Trigger class

---
 CHANGELOG.md             |  7 +++++
 docs/sphinx/api.rst      | 11 +++++++
 src/lvmopstools/utils.py | 64 ++++++++++++++++++++++++++++++++++++++++
 tests/test_utils.py      | 49 +++++++++++++++++++++++++++++-
 4 files changed, 130 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index f9884e2..0fd2a28 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,5 +1,12 @@
 # Changelog
 
+## Next version
+
+### 🚀 New
+
+* Added `Trigger` class.
+
+
 ## 0.5.0 - December 21, 2024
 
 ### 💥 Breaking changes
diff --git a/docs/sphinx/api.rst b/docs/sphinx/api.rst
index 938047c..77b6d3c 100644
--- a/docs/sphinx/api.rst
+++ b/docs/sphinx/api.rst
@@ -54,6 +54,17 @@ Socket
 .. autoclass:: lvmopstools.socket.AsyncSocketHandler
    :members:
 
+Utils
+-----
+
+.. autofunction:: lvmopstools.utils.get_amqp_client
+.. autofunction:: lvmopstools.utils.get_exception_data
+.. autofunction:: lvmopstools.utils.stop_event_loop
+.. autofunction:: lvmopstools.utils.with_timeout
+.. autofunction:: lvmopstools.utils.is_notebook
+.. autoclass:: lvmopstools.utils.Trigger
+
+
 Weather
 -------
 
diff --git a/src/lvmopstools/utils.py b/src/lvmopstools/utils.py
index 6cb6d5f..20e3406 100644
--- a/src/lvmopstools/utils.py
+++ b/src/lvmopstools/utils.py
@@ -9,6 +9,7 @@
 from __future__ import annotations
 
 import asyncio
+import time
 
 from typing import Any, Coroutine, TypeVar
 
@@ -21,6 +22,7 @@
     "stop_event_loop",
     "with_timeout",
     "is_notebook",
+    "Trigger",
 ]
 
 
@@ -138,3 +140,65 @@ async def with_timeout(
     except asyncio.TimeoutError:
         if raise_on_timeout:
             raise asyncio.TimeoutError(f"Timed out after {timeout} seconds.")
+
+
+class Trigger:
+    """A trigger that can be set and reset and accepts setting thresholds.
+
+    This class is essentially just a flag that can take true/false values, but
+    triggering the true value can be delayed by time or number or triggers.
+
+    Parameters
+    ----------
+    n
+        The number of times the instance needs to be set before it is triggered.
+    delay
+        The delay in seconds before the instance is triggered. This is counted from
+        the first time the instance is set, and is reset if the instance is reset.
+        If ``n_triggers`` is greater than 1, both conditions must be met for the
+        instance to be triggered.
+
+    """
+
+    def __init__(self, n: int = 1, delay: float = 0):
+        self.n = n
+        self.delay = delay
+
+        self._triggered = False
+        self._first_trigger: float | None = None
+        self._n_sets: int = 0
+
+    def _check(self):
+        """Check the trigger conditions and update the internal state."""
+
+        now = time.time()
+        if (
+            self._n_sets >= self.n
+            and self._first_trigger is not None
+            and now - self._first_trigger >= self.delay
+        ):
+            self._triggered = True
+
+    def set(self):
+        """Sets the trigger."""
+
+        if self._triggered:
+            return
+
+        self._n_sets += 1
+        self._first_trigger = self._first_trigger or time.time()
+        self._check()
+
+    def reset(self):
+        """Resets the trigger."""
+
+        self._first_trigger = None
+        self._n_sets = 0
+        self._triggered = False
+
+    def is_set(self):
+        """Returns :obj:`True` if the trigger is set."""
+
+        self._check()
+
+        return self._triggered
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 8cea86e..ee0445d 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -14,7 +14,7 @@
 import pytest_mock
 
 import lvmopstools.utils
-from lvmopstools.utils import is_notebook, with_timeout
+from lvmopstools.utils import Trigger, is_notebook, with_timeout
 
 
 async def _timeout(delay: float):
@@ -69,3 +69,50 @@ async def test_is_notebook_name_Error(mocker: pytest_mock.MockerFixture):
     )
 
     assert not is_notebook()
+
+
+async def test_trigger_n_sets():
+    trigger = Trigger(n=3)
+
+    assert not trigger.is_set()
+
+    trigger.set()
+    assert not trigger.is_set()
+
+    trigger.set()
+    assert not trigger.is_set()
+
+    trigger.set()
+    assert trigger.is_set()
+
+
+async def test_trigger_n_sets_delay():
+    trigger = Trigger(n=2, delay=0.25)
+
+    trigger.set()
+    assert not trigger.is_set()
+
+    trigger.set()
+    assert not trigger.is_set()
+
+    await asyncio.sleep(0.3)
+    assert trigger.is_set()
+
+    trigger.set()
+    assert trigger.is_set()
+
+
+async def test_trigger_reset():
+    trigger = Trigger(n=2)
+
+    trigger.set()
+    trigger.set()
+    assert trigger.is_set()
+
+    trigger.reset()
+    assert not trigger.is_set()
+
+    trigger.set()
+    trigger.reset()
+    trigger.set()
+    assert not trigger.is_set()