From c3e9690280df25a916a89f0b744dd5fc431bec14 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Thu, 17 Feb 2022 14:25:03 -0500
Subject: [PATCH] Agent: Add request_cache decorator

---
 monkey/infection_monkey/utils/decorators.py   | 43 ++++++++++
 .../infection_monkey/utils/test_decorators.py | 78 +++++++++++++++++++
 2 files changed, 121 insertions(+)
 create mode 100644 monkey/infection_monkey/utils/decorators.py
 create mode 100644 monkey/tests/unit_tests/infection_monkey/utils/test_decorators.py

diff --git a/monkey/infection_monkey/utils/decorators.py b/monkey/infection_monkey/utils/decorators.py
new file mode 100644
index 000000000..7a93a7c7a
--- /dev/null
+++ b/monkey/infection_monkey/utils/decorators.py
@@ -0,0 +1,43 @@
+from functools import wraps
+
+from .timer import Timer
+
+
+def request_cache(ttl: float):
+    """
+    This is a decorator that allows a single response of a function to be cached with an expiration
+    time (TTL). The first call to the function is executed and the response is cached. Subsequent
+    calls to the function result in the cached value being returned until the TTL elapses. Once the
+    TTL elapses, the cache is considered stale and the decorated function will be called, its
+    response cached, and the TTL reset.
+
+    An example usage of this decorator is to wrap a function that makes frequent slow calls to an
+    external resource, such as an HTTP request to a remote endpoint. If the most up-to-date
+    information is not need, this decorator provides a simple way to cache the response for a
+    certain amount of time.
+
+    Example:
+        @request_cache(600)
+        def raining_outside():
+            return requests.get(f"https://weather.service.api/check_for_rain/{MY_ZIP_CODE}")
+
+    :param ttl: The time-to-live in seconds for the cached return value
+    :return: The return value of the decorated function, or the cached return value if the TTL has
+             not elapsed.
+    """
+
+    def decorator(fn):
+        @wraps(fn)
+        def wrapper(*args, **kwargs):
+            if wrapper.timer.is_expired():
+                wrapper.cached_value = fn(*args, **kwargs)
+                wrapper.timer.set(ttl)
+
+            return wrapper.cached_value
+
+        wrapper.cached_value = None
+        wrapper.timer = Timer()
+
+        return wrapper
+
+    return decorator
diff --git a/monkey/tests/unit_tests/infection_monkey/utils/test_decorators.py b/monkey/tests/unit_tests/infection_monkey/utils/test_decorators.py
new file mode 100644
index 000000000..30af5837a
--- /dev/null
+++ b/monkey/tests/unit_tests/infection_monkey/utils/test_decorators.py
@@ -0,0 +1,78 @@
+import time
+from unittest.mock import MagicMock
+
+import pytest
+
+from infection_monkey.utils.decorators import request_cache
+from infection_monkey.utils.timer import Timer
+
+
+class MockTimer(Timer):
+    def __init__(self):
+        self._time_remaining = 0
+        self._set_time = 0
+
+    def set(self, timeout_sec: float):
+        self._time_remaining = timeout_sec
+        self._set_time = timeout_sec
+
+    def set_expired(self):
+        self._time_remaining = 0
+
+    @property
+    def time_remaining(self) -> float:
+        return self._time_remaining
+
+    def reset(self):
+        """
+        Reset the timer without changing the timeout
+        """
+        self._time_remaining = self._set_time
+
+
+class MockTimerFactory:
+    def __init__(self):
+        self._instance = None
+
+    def __call__(self):
+        if self._instance is None:
+            mt = MockTimer()
+            self._instance = mt
+
+        return self._instance
+
+    def reset(self):
+        self._instance = None
+
+
+mock_timer_factory = MockTimerFactory()
+
+
+@pytest.fixture
+def mock_timer(monkeypatch):
+    mock_timer_factory.reset
+
+    monkeypatch.setattr("infection_monkey.utils.decorators.Timer", mock_timer_factory)
+
+    return mock_timer_factory()
+
+
+def test_request_cache(mock_timer):
+    mock_request = MagicMock(side_effect=lambda: time.time())
+
+    @request_cache(10)
+    def make_request():
+        return mock_request()
+
+    t1 = make_request()
+    t2 = make_request()
+
+    assert t1 == t2
+
+    mock_timer.set_expired()
+
+    t3 = make_request()
+    t4 = make_request()
+
+    assert t3 != t1
+    assert t3 == t4