diff --git a/conftest.py b/conftest.py index 397f620..c41a328 100644 --- a/conftest.py +++ b/conftest.py @@ -15,6 +15,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter from opentelemetry.trace import set_tracer_provider +from appsignal import probes from appsignal.agent import agent from appsignal.opentelemetry import METRICS_PREFERRED_TEMPORALITY @@ -87,13 +88,11 @@ def remove_logging_handlers_after_tests(): @pytest.fixture(scope="function", autouse=True) -def remove_probes_after_tests(): +def stop_and_clear_probes_after_tests(): yield - from appsignal.probes import _probes, _probe_states - - _probes.clear() - _probe_states.clear() + probes.stop() + probes.clear() @pytest.fixture(scope="function", autouse=True) diff --git a/src/appsignal/probes.py b/src/appsignal/probes.py index 615cf60..7e6ded7 100644 --- a/src/appsignal/probes.py +++ b/src/appsignal/probes.py @@ -2,9 +2,9 @@ import logging from inspect import signature -from threading import Lock, Thread -from time import gmtime, sleep -from typing import Any, Callable, NoReturn, Optional, TypeVar, Union, cast +from threading import Event, Lock, Thread +from time import gmtime +from typing import Any, Callable, Optional, TypeVar, Union, cast T = TypeVar("T") @@ -15,6 +15,7 @@ _probe_states: dict[str, Any] = {} _lock: Lock = Lock() _thread: Thread | None = None +_stop_event: Event = Event() def start() -> None: @@ -24,12 +25,15 @@ def start() -> None: _thread.start() -def _minutely_loop() -> NoReturn: - sleep(_initial_wait_time()) +def _minutely_loop() -> None: + wait_time = _initial_wait_time() while True: + if _stop_event.wait(timeout=wait_time): + break + _run_probes() - sleep(_wait_time()) + wait_time = _wait_time() def _run_probes() -> None: @@ -88,3 +92,18 @@ def unregister(name: str) -> None: del _probes[name] if name in _probe_states: del _probe_states[name] + + +def stop() -> None: + global _thread + if _thread is not None: + _stop_event.set() + _thread.join() + _thread = None + _stop_event.clear() + + +def clear() -> None: + with _lock: + _probes.clear() + _probe_states.clear() diff --git a/tests/test_probes.py b/tests/test_probes.py index 94ec4a3..22300b6 100644 --- a/tests/test_probes.py +++ b/tests/test_probes.py @@ -1,6 +1,15 @@ +from time import sleep from typing import Any, Callable, cast -from appsignal.probes import _probe_states, _probes, _run_probes, register, unregister +from appsignal.probes import ( + _probe_states, + _probes, + _run_probes, + register, + start, + stop, + unregister, +) def test_register(mocker): @@ -32,7 +41,7 @@ def test_register_with_state(mocker): probe.assert_called_with("state") -def test_register_signatures(mocker): +def test_register_signatures(): # `mocker.Mock` is not used here because we want to test against # specific function signatures. @@ -107,3 +116,23 @@ def test_unregister(mocker): assert "probe_name" not in _probes assert "probe_name" not in _probe_states probe.assert_called_once_with(None) + + +def test_start_stop(mocker): + mocker.patch("appsignal.probes._initial_wait_time").return_value = 0.001 + mocker.patch("appsignal.probes._wait_time").return_value = 0.001 + + probe = mocker.Mock() + register("probe_name", probe) + start() + + sleep(0.05) + + probe.assert_called() + + stop() + call_count = probe.call_count + + sleep(0.05) + + assert probe.call_count == call_count