Expand and refactor state tracking (#70)

* Add mock Zabbix API autouse fixture

* Rewrite process state tracking

Uses discrete State class instead of passing around a dict.

* Remove unused import

* Refactor source collector exception msg chaining

* Add test for State.asdict with State.set_error

tests/conftest.py

@@ -1,8 +1,10 @@
 import multiprocessing
 import os
 from pathlib import Path
-from typing import Iterable
+from typing import Iterable, Type
+from unittest.mock import MagicMock
 import pytest
+from unittest import mock
 
 
 @pytest.fixture(scope="function")
@@ -129,4 +131,33 @@ def setup_multiprocessing_start_method() -> None:
     # On MacOS we have to set the start mode to fork
     # when using multiprocessing-logging
     if os.uname == "Darwin":
-        multiprocessing.set_start_method("fork", force=True)
+        multiprocessing.set_start_method("fork", force=True)
+
+
+class PicklableMock(MagicMock):
+    def __reduce__(self):
+        return (MagicMock, ())
+
+
+class MockZabbixAPI(PicklableMock):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.apiinfo = PicklableMock()
+        self.apiinfo.version = PicklableMock(return_value="5.0.0")
+        self.login = PicklableMock()
+
+
+# NOTE: if doing integration testing in the future, the definitions of these
+# fixtures should be dependent on some env var, which enables/disables mocking
+
+
+@pytest.fixture(autouse=True)
+def mock_zabbix_api() -> Iterable[Type[MockZabbixAPI]]:
+    with mock.patch("pyzabbix.ZabbixAPI", new=MockZabbixAPI) as api_mock:
+        yield api_mock
+
+
+@pytest.fixture()
+def mock_psycopg2_connect() -> Iterable[PicklableMock]:
+    with mock.patch("psycopg2.connect", PicklableMock()) as psycopg_mock:
+        yield psycopg_mock

tests/test_processing/test_sourcecollectorprocess.py

@@ -6,6 +6,7 @@
 
 from zabbix_auto_config.processing import SourceCollectorProcess
 from zabbix_auto_config.models import Host, SourceCollectorSettings
+from zabbix_auto_config.state import get_manager
 
 
 class SourceCollector:
@@ -22,7 +23,7 @@ def collect(*args, **kwargs) -> List[Host]:
 def test_source_collector_process():
     process = SourceCollectorProcess(
         name="test-source",
-        state=multiprocessing.Manager().dict(),
+        state=get_manager().State(),
         module=SourceCollector,
         config=SourceCollectorSettings(
             module_name="source_collector",
@@ -40,7 +41,7 @@ def test_source_collector_process():
         hosts = process.source_hosts_queue.get()
         assert len(hosts["hosts"]) == 2
         assert hosts["hosts"][0].hostname == "foo.example.com"
-        assert process.state["ok"] is True
+        assert process.state.ok is True
     finally:
         process.stop_event.set()
         process.join(timeout=0.01)
@@ -57,7 +58,7 @@ def collect(*args, **kwargs) -> List[Host]:
 def test_source_collector_disable_on_failure():
     process = SourceCollectorProcess(
         name="test-source",
-        state=multiprocessing.Manager().dict(),
+        state=get_manager().State(),
         module=FaultySourceCollector,
         config=SourceCollectorSettings(
             module_name="faulty_source_collector",
@@ -73,9 +74,9 @@ def test_source_collector_disable_on_failure():
     # Start process and wait until it fails
     try:
         process.start()
-        while process.state["ok"] is True:
+        while process.state.ok is True:
             time.sleep(0.01)
-        assert process.state["ok"] is False
+        assert process.state.ok is False
         assert process.source_hosts_queue.empty() is True
         process.stop_event.set()
     finally:

tests/test_processing/test_zabbixupdater.py

@@ -1,39 +1,47 @@
-import multiprocessing
 from pathlib import Path
 import time
-from unittest.mock import MagicMock, patch, Mock
+from unittest.mock import patch
 import pytest
 import requests
+
+from ..conftest import MockZabbixAPI, PicklableMock
 from zabbix_auto_config import exceptions
 
 from zabbix_auto_config.models import ZabbixSettings
 from zabbix_auto_config.processing import ZabbixUpdater
+from zabbix_auto_config.state import get_manager
 
 
 def raises_connect_timeout(*args, **kwargs):
     raise requests.exceptions.ConnectTimeout("connect timeout")
 
 
+# We have to set the side effect in the constructor
+class TimeoutAPI(MockZabbixAPI):
+    def __init__(self, *args, **kwargs) -> None:
+        super().__init__(*args, **kwargs)
+        self.login = PicklableMock(
+            side_effect=requests.exceptions.ConnectTimeout("connect timeout")
+        )
+
+
 @pytest.mark.timeout(10)
-@patch("psycopg2.connect", MagicMock())  # throwaway mock
-def test_zabbixupdater_connect_timeout():
+@patch("pyzabbix.ZabbixAPI", TimeoutAPI())  # mock with timeout on login
+def test_zabbixupdater_connect_timeout(mock_psycopg2_connect):
     with pytest.raises(exceptions.ZACException) as exc_info:
-        with patch(
-            "pyzabbix.ZabbixAPI.login", new_callable=lambda: raises_connect_timeout
-        ):
-            ZabbixUpdater(
-                name="connect-timeout",
-                db_uri="",
-                state=multiprocessing.Manager().dict(),
-                zabbix_config=ZabbixSettings(
-                    map_dir="",
-                    url="",
-                    username="",
-                    password="",
-                    dryrun=False,
-                    timeout=1,
-                ),
-            )
+        ZabbixUpdater(
+            name="connect-timeout",
+            db_uri="",
+            state=get_manager().State(),
+            zabbix_config=ZabbixSettings(
+                map_dir="",
+                url="",
+                username="",
+                password="",
+                dryrun=False,
+                timeout=1,
+            ),
+        )
     assert "connect timeout" in exc_info.exconly()
 
 
@@ -42,15 +50,8 @@ def do_update(self):
         raise requests.exceptions.ReadTimeout("read timeout")
 
 
-class PickableMock(MagicMock):
-    def __reduce__(self):
-        return (MagicMock, ())
-
-
 @pytest.mark.timeout(5)
-@patch("psycopg2.connect", PickableMock())
-@patch("pyzabbix.ZabbixAPI", PickableMock())
-def test_zabbixupdater_read_timeout(tmp_path: Path):
+def test_zabbixupdater_read_timeout(tmp_path: Path, mock_psycopg2_connect):
     # TODO: use mapping file fixtures from #67
     map_dir = tmp_path / "maps"
     map_dir.mkdir()
@@ -61,7 +62,7 @@ def test_zabbixupdater_read_timeout(tmp_path: Path):
     process = TimeoutUpdater(
         name="read-timeout",
         db_uri="",
-        state=multiprocessing.Manager().dict(),
+        state=get_manager().State(),
         zabbix_config=ZabbixSettings(
             map_dir=str(map_dir),
             url="",
@@ -75,9 +76,11 @@ def test_zabbixupdater_read_timeout(tmp_path: Path):
     # Start the process and wait for it to be marked as unhealthy
     try:
         process.start()
-        while process.state["ok"] is True:
+        while process.state.ok is True:
             time.sleep(0.1)
-        assert process.state["ok"] is False
+        assert process.state.ok is False
+        assert process.state.error_type == "ReadTimeout"
+        assert process.state.error_count == 1
         process.stop_event.set()
     finally:
         process.join(timeout=0.01)

tests/test_state.py

@@ -0,0 +1,178 @@
+import datetime
+import time
+
+import pytest
+from zabbix_auto_config.exceptions import ZACException
+from zabbix_auto_config.processing import BaseProcess
+
+from zabbix_auto_config.state import State, get_manager, StateProxy
+
+
+def test_manager_state():
+    manager = get_manager()
+    state = manager.State()
+    assert isinstance(state, StateProxy)
+    # Test defaults
+    assert state.ok is True
+    assert state.error is None
+    assert state.error_type is None
+    assert state.error_count == 0
+    assert state.error_time is None
+
+
+@pytest.mark.parametrize("use_manager", [True, False])
+@pytest.mark.parametrize("with_error", [True, False])
+def test_state_set_ok(use_manager: bool, with_error: bool):
+    if use_manager:
+        state = get_manager().State()
+    else:
+        state = State()
+
+    # Give state object some error state values
+    if with_error:
+        state.error = "Error"
+        state.error_type = "ErrorType"
+        state.error_count = 1
+        state.error_time = datetime.datetime(2021, 1, 1, 0, 0, 0).timestamp()
+
+    state.set_ok()
+    assert state.ok is True
+    assert state.error is None
+    assert state.error_type is None
+    assert state.error_time is None
+    if with_error:
+        assert state.error_count == 1
+    else:
+        assert state.error_count == 0
+
+
+# Use a subclass of Exception so that we can test that
+# the error type is set correctly
+# Also needs to be in the global scope to be pickleable
+class TimeoutError(Exception):
+    pass
+
+
+@pytest.mark.parametrize("use_manager", [True, False])
+def test_state_set_error(use_manager: bool):
+    if use_manager:
+        state = get_manager().State()
+    else:
+        state = State()
+
+    # Sanity test of defaults
+    assert state.ok is True
+    assert state.error is None
+    assert state.error_type is None
+    assert state.error_time is None
+    assert state.error_count == 0
+
+    time.sleep(0.01)  # to ensure later timestamps are greater
+    e = TimeoutError("Test error")
+    state.set_error(e)
+    assert state.ok is False
+    assert state.error == "Test error"
+    assert state.error_type == "TimeoutError"
+    assert state.error_time < time.time()
+    assert state.error_count == 1
+
+    # Set the error again to check count and time are updated
+    prev_time = float(state.error_time)
+    state.set_error(e)
+    assert state.error_count == 2
+    assert state.error_time > prev_time
+
+
+class ZACExceptionProcess(BaseProcess):
+    def work(self) -> None:
+        raise ZACException("Test error")
+
+
+@pytest.mark.timeout(10)
+def test_state_in_other_process() -> None:
+    state = get_manager().State()
+    process = ZACExceptionProcess(
+        name="test",
+        state=state,
+    )
+
+    process.start()
+    try:
+        while process.state.ok:
+            time.sleep(0.01)
+        process.stop_event.set()
+    finally:
+        # stop process to prevent errors from accumulating
+        process.join(timeout=0.01)
+
+    assert process.state.ok is False
+    assert process.state.error_type == "ZACException"
+    assert process.state.error_count == 1
+    assert process.state is state
+
+    # Test that multiple state proxies do not refer to the same
+    # underlying State object
+    state2 = get_manager().State()
+    assert state2.ok is True
+    assert state2 is not state
+    # This process will not fail and thus will set its state to OK
+    process2 = BaseProcess(
+        name="test",
+        state=state2,
+    )
+
+    # Start and stop process, then check state
+    try:
+        process2.start()
+        process2.stop_event.set()
+    finally:
+        process2.join(timeout=1)
+    assert process2.state.ok is True
+    assert process2.state.asdict() == state2.asdict()
+    assert process2.state.asdict() != process.state.asdict()
+    assert process2.state is not process.state
+
+
+@pytest.mark.parametrize("use_manager", [True, False])
+def test_state_asdict_ok(use_manager: bool) -> None:
+    if use_manager:
+        state = get_manager().State()
+    else:
+        state = State()
+    state.set_ok()
+    assert state.asdict() == {
+        "ok": True,
+        "error": None,
+        "error_type": None,
+        "error_count": 0,
+        "error_time": None,
+    }
+
+
+class CustomException(Exception):
+    pass
+
+
+@pytest.mark.parametrize("use_manager", [True, False])
+def test_state_asdict_error(use_manager: bool) -> None:
+    if use_manager:
+        state = get_manager().State()
+    else:
+        state = State()
+
+    # Mocking datetime in subprocesses is a bit of a chore, so we just
+    # check that the error_time is a timestamp value within a given range
+    pre = datetime.datetime.now().timestamp()
+    state.set_error(CustomException("Test error"))
+    post = datetime.datetime.now().timestamp()
+    d = state.asdict()
+
+    assert post >= d["error_time"] >= pre
+    d.pop("error_time")
+
+    assert d == {
+        "ok": False,
+        "error": "Test error",
+        "error_type": "CustomException",
+        "error_count": 1,
+    }