From e981ead1500e58892a9317c07e456f13c01dc708 Mon Sep 17 00:00:00 2001
From: Mike Salvatore <mike.s.salvatore@gmail.com>
Date: Wed, 9 Feb 2022 18:18:15 -0500
Subject: [PATCH] Agent: Add new time_remaining() method to Timer

---
 monkey/infection_monkey/utils/timer.py        | 13 +++++++++-
 .../infection_monkey/utils/test_timer.py      | 25 +++++++++++++++++++
 2 files changed, 37 insertions(+), 1 deletion(-)

diff --git a/monkey/infection_monkey/utils/timer.py b/monkey/infection_monkey/utils/timer.py
index 366a10b20..597eb6020 100644
--- a/monkey/infection_monkey/utils/timer.py
+++ b/monkey/infection_monkey/utils/timer.py
@@ -25,7 +25,18 @@ class Timer:
                  TIMEOUT_SEC, False otherwise
         :rtype: bool
         """
-        return (time.time() - self._start_time) >= self._timeout_sec
+        return self.time_remaining == 0
+
+    @property
+    def time_remaining(self) -> float:
+        """
+        Return the amount of time remaining until the timer expires.
+        :return: The number of seconds until the timer expires. If the timer is expired, this
+                 function returns 0 (it will never return a negative number).
+        :rtype: float
+        """
+        time_remaining = self._timeout_sec - (time.time() - self._start_time)
+        return time_remaining if time_remaining > 0 else 0
 
     def reset(self):
         """
diff --git a/monkey/tests/unit_tests/infection_monkey/utils/test_timer.py b/monkey/tests/unit_tests/infection_monkey/utils/test_timer.py
index 5359b8c79..b5291cc0e 100644
--- a/monkey/tests/unit_tests/infection_monkey/utils/test_timer.py
+++ b/monkey/tests/unit_tests/infection_monkey/utils/test_timer.py
@@ -67,3 +67,28 @@ def test_timer_reset(start_time, set_current_time, timeout):
 
     set_current_time(start_time + (2 * timeout))
     assert t.is_expired()
+
+
+def test_time_remaining(start_time, set_current_time):
+    timeout = 5
+
+    t = Timer()
+    t.set(timeout)
+
+    assert t.time_remaining == timeout
+
+    set_current_time(start_time + 2)
+    assert t.time_remaining == 3
+
+
+def test_time_remaining_is_zero(start_time, set_current_time):
+    timeout = 5
+
+    t = Timer()
+    t.set(timeout)
+
+    set_current_time(start_time + timeout)
+    assert t.time_remaining == 0
+
+    set_current_time(start_time + (2 * timeout))
+    assert t.time_remaining == 0