PWMProbes: Unit tests for heartbeat thread

IBM#38
Signed-off-by: Gabe Goodhart <ghart@us.ibm.com>

oper8/test_helpers/pwm_helpers.py

@@ -2,12 +2,14 @@
 Utils and common classes for the python watch manager tests
 """
 # Standard
+from datetime import datetime
 from multiprocessing.connection import Connection
 from queue import Queue
 from threading import Event
 from uuid import uuid4
 import multiprocessing
 import random
+import tempfile
 import time
 
 # Third Party
@@ -31,6 +33,7 @@
 from oper8.watch_manager.python_watch_manager.leader_election.life import (
     LeaderForLifeManager,
 )
+from oper8.watch_manager.python_watch_manager.threads.heartbeat import HeartbeatThread
 from oper8.watch_manager.python_watch_manager.threads.reconcile import ReconcileThread
 from oper8.watch_manager.python_watch_manager.threads.timer import TimerThread
 from oper8.watch_manager.python_watch_manager.utils.types import (
@@ -79,6 +82,10 @@ class MockedTimerThread(TimerThread):
     _disable_singleton = True
 
 
+class MockedHeartbeatThread(HeartbeatThread):
+    _disable_singleton = True
+
+
 class MockedReconcileThread(ReconcileThread):
     """Subclass of ReconcileThread that mocks the subprocess. This was more
     reliable than using unittest.mock"""
@@ -212,3 +219,17 @@ def mocked_create_and_start_entrypoint(
     time.sleep(wait_time)
     for message in returned_messages or []:
         result_pipe.send(message)
+
+
+def read_heartbeat_file(hb_file: str) -> datetime:
+    """Parse a heartbeat file into a datetime"""
+    with open(hb_file) as handle:
+        hb_str = handle.read()
+
+    return datetime.strptime(hb_str, HeartbeatThread._DATE_FORMAT)
+
+
+@pytest.fixture
+def heartbeat_file():
+    with tempfile.NamedTemporaryFile() as tmp_file:
+        yield tmp_file.name

tests/watch_manager/python_watch_manager/threads/test_heartbeat.py

@@ -3,29 +3,89 @@
 """
 # Standard
 from datetime import datetime, timedelta
-import tempfile
-import time
-
-# Third Party
-import pytest
+from unittest import mock
 
 # Local
-from oper8.exceptions import ConfigError
-from oper8.watch_manager.python_watch_manager.threads import HeartbeatThread
+from oper8.test_helpers.pwm_helpers import (
+    MockedHeartbeatThread,
+    heartbeat_file,
+    read_heartbeat_file,
+)
+
+## Helpers #####################################################################
+
+
+class FailOnceOpen:
+    def __init__(self, fail_on: int = 1):
+        self.call_num = 0
+        self.fail_on = fail_on
+        self._real_open = open
+
+    def __call__(self, *args, **kwargs):
+        self.call_num += 1
+        if self.call_num == self.fail_on:
+            print(f"Raising on call {self.call_num}")
+            raise OSError("Yikes")
+        print(f"Returning from call {self.call_num}")
+        return self._real_open(*args, **kwargs)
+
+
+## Tests #######################################################################
+
+
+def test_heartbeat_happy_path(heartbeat_file):
+    """Make sure the heartbeat initializes correctly"""
+    hb = MockedHeartbeatThread(heartbeat_file, "1s")
 
+    # Heartbeat not run until started
+    with open(heartbeat_file) as handle:
+        assert not handle.read()
 
-class NonSingletonHeartbeatThread(HeartbeatThread):
-    _disable_singleton = True
+    # Start and stop the thread to trigger the first heartbeat only
+    hb.start_thread()
+    hb.wait_for_beat()
+    hb.stop_thread()
 
+    # Make sure the heartbeat is "current"
+    assert read_heartbeat_file(heartbeat_file) > (datetime.now() - timedelta(seconds=5))
 
-def test_simple_heartbeat():
-    with tempfile.NamedTemporaryFile() as heartbeat_file:
-        hb = NonSingletonHeartbeatThread(heartbeat_file.name, "1s")
-        hb.start()
-        time.sleep(1)
+
+def test_heartbeat_ongoing(heartbeat_file):
+    """Make sure that the heartbeat continues to beat in an ongoing way"""
+    hb = MockedHeartbeatThread(heartbeat_file, "1s")
+
+    # Start the thread and read the first one
+    hb.start_thread()
+    hb.wait_for_beat()
+    first_hb = read_heartbeat_file(heartbeat_file)
+
+    # Wait a bit and read again
+    hb.wait_for_beat()
+    hb.stop_thread()
+    later_hb = read_heartbeat_file(heartbeat_file)
+    assert later_hb > first_hb
+
+
+def test_heartbeat_with_exception(heartbeat_file):
+    """Make sure that a sporadic failure does not terminate the heartbeat"""
+    # Mock so that the third call to open will raise. This correlates with the
+    # second heartbeat since we read the file using open after each heartbeat
+    with mock.patch("builtins.open", new=FailOnceOpen(3)):
+        hb = MockedHeartbeatThread(heartbeat_file, "1s")
+        hb.start_thread()
+
+        # The first beat succeeds
+        hb.wait_for_beat()
+        first_hb = read_heartbeat_file(heartbeat_file)
+
+        # The first beat raises, but doesn't cause any problems
+        hb.wait_for_beat()
+        second_hb = read_heartbeat_file(heartbeat_file)
+
+        # The third beat succeeds
+        hb.wait_for_beat()
+        third_hb = read_heartbeat_file(heartbeat_file)
         hb.stop_thread()
-        hb.join()
-        with open(heartbeat_file.name) as handle:
-            hb_str = handle.read()
-        parsed = datetime.strptime(HeartbeatThread._DATE_FORMAT, hb_str)
-        assert parsed > (datetime.now() - timedelta(seconds=5))
+
+        assert first_hb == second_hb
+        assert third_hb > first_hb