Fix flaky timed tests

README.rst

@@ -22,7 +22,7 @@ Infrastructure Continuous Ordinance Mapping for Planning and Siting Systems (INF
 .. |SWR| image:: https://img.shields.io/badge/SWR--25--62_-blue?label=NREL
     :alt: Static Badge
 
-.. |Zenodo| image:: https://zenodo.org/badge/892830182.svg
+.. |Zenodo| image:: https://zenodo.org/badge/DOI/10.5281/zenodo.17173409.svg
     :target: https://doi.org/10.5281/zenodo.17173409
 
 .. inclusion-intro

tests/python/conftest.py

@@ -1,5 +1,6 @@
 """Fixtures for use across all tests"""
 
+import time
 import asyncio
 from pathlib import Path
 
@@ -125,3 +126,26 @@ def _get_response(
         )
 
     return _get_response
+
+
+@pytest.fixture
+def patched_clock(monkeypatch):
+    """Fixture to patch time.monotonic with a deterministic clock"""
+
+    class FakeClock:
+        """Deterministic replacement for ``time.monotonic`` in tests"""
+
+        def __init__(self, start=0.0):
+            self._now = start
+
+        def __call__(self):
+            return self._now
+
+        def advance(self, seconds):
+            self._now += seconds
+            return self._now
+
+    fake_clock = FakeClock()
+    monkeypatch.setattr(time, "monotonic", fake_clock)
+
+    return fake_clock

tests/python/integration/test_integrated.py

@@ -22,6 +22,7 @@
 from compass.services.provider import RunningAsyncServices
 from compass.utilities.enums import LLMUsageCategory
 from compass.utilities.logs import LocationFileLog, LogListener
+from elm.utilities import retry as retry_module
 
 
 class MockResponse:
@@ -60,14 +61,27 @@ async def patched_get_html(url, *args, **kwargs):  # noqa: RUF029
 
 
 @pytest.mark.asyncio
-async def test_openai_query(sample_openai_response, monkeypatch):
+async def test_openai_query(
+    sample_openai_response, monkeypatch, patched_clock
+):
     """Test querying OpenAI while tracking limits and usage"""
 
     start_time = None
     elapsed_times = []
     time_limit = 0.1
     sleep_mult = 1.2
 
+    original_sleep = asyncio.sleep
+
+    async def fake_sleep(delay, result=None):
+        patched_clock.advance(delay)
+        await original_sleep(0)
+        return result
+
+    monkeypatch.setattr(asyncio, "sleep", fake_sleep)
+    monkeypatch.setattr(retry_module.asyncio, "sleep", fake_sleep)
+    monkeypatch.setattr(retry_module.random, "random", lambda: 0.0)
+
     async def _test_response(*args, **kwargs):  # noqa: RUF029
         time_elapsed = time.monotonic() - start_time
         elapsed_times.append(time_elapsed)
@@ -104,6 +118,7 @@ async def _test_response(*args, **kwargs):  # noqa: RUF029
     async with RunningAsyncServices([openai_service]):
         start_time = time.monotonic()
         message = await openai_service.call(usage_tracker=usage_tracker)
+        patched_clock.advance(time_limit * 3)
         message2 = await openai_service.call()
 
         assert openai_service.rate_tracker.total == 13
@@ -124,19 +139,20 @@ async def _test_response(*args, **kwargs):  # noqa: RUF029
             }
         }
 
-        time.sleep(time_limit * sleep_mult)
+        patched_clock.advance(time_limit * sleep_mult)
         assert openai_service.rate_tracker.total == 0
 
         start_time = time.monotonic() - time_limit - 1
         await openai_service.call()
+        patched_clock.advance(time_limit * sleep_mult * 0.8 + 0.01)
         await openai_service.call()
         assert len(elapsed_times) == 5
         assert elapsed_times[-2] - time_limit - 1 < 1
         assert (
             elapsed_times[-1] - time_limit - 1 > time_limit * sleep_mult * 0.8
         )
 
-        time.sleep(time_limit * sleep_mult)
+        patched_clock.advance(time_limit * sleep_mult)
         start_time = time.monotonic() - time_limit - 1
         assert openai_service.rate_tracker.total == 0
 

tests/python/unit/services/test_services_base.py

@@ -1,17 +1,14 @@
 """Test Ordinances Base Services"""
 
-import time
 from pathlib import Path
 
 import pytest
-from flaky import flaky
 
 from compass.services.base import LLMService
 from compass.services.usage import TimeBoundedUsageTracker
 
 
-@flaky(max_runs=10, min_passes=1)
-def test_base_llm_limited_service():
+def test_base_llm_limited_service(patched_clock):
     """Test base implementation of `LLMService` class"""
 
     class TestService(LLMService):
@@ -29,9 +26,10 @@ async def process(self, *args, **kwargs):
     assert service.can_process
     service.rate_tracker.add(50)
     assert service.can_process
+    patched_clock.advance(0.01)
     service.rate_tracker.add(75)
     assert not service.can_process
-    time.sleep(0.1)
+    patched_clock.advance(0.099)
     assert service.can_process
 
 

tests/python/unit/services/test_services_usage.py

@@ -1,10 +1,8 @@
 """Test COMPASS Ordinance service usage functions and classes"""
 
-import time
 from pathlib import Path
 
 import pytest
-from flaky import flaky
 
 from compass.services.usage import (
     TimedEntry,
@@ -23,15 +21,15 @@ def _sample_response_parser(current_usage, response):
     return current_usage
 
 
-def test_timed_entry():
+def test_timed_entry(patched_clock):
     """Test `TimedEntry` class"""
 
     a = TimedEntry(100)
-    assert a <= time.monotonic()
+    assert a <= patched_clock()
 
-    time.sleep(0.2)
-    sample_time = time.monotonic()
-    time.sleep(0.2)
+    patched_clock.advance(0.2)
+    sample_time = patched_clock()
+    patched_clock.advance(0.2)
     b = TimedEntry(10000)
     assert b > sample_time
     assert a < sample_time
@@ -40,20 +38,19 @@ def test_timed_entry():
     assert b.value == 10000
 
 
-@flaky(max_runs=10, min_passes=1)
-def test_time_bounded_usage_tracker():
+def test_time_bounded_usage_tracker(patched_clock):
     """Test the `TimeBoundedUsageTracker` class"""
 
     tracker = TimeBoundedUsageTracker(max_seconds=0.2)
     assert tracker.total == 0
     tracker.add(500)
     assert tracker.total == 500
-    time.sleep(0.1)
+    patched_clock.advance(0.1)
     tracker.add(200)
     assert tracker.total == 700
-    time.sleep(0.1)
+    patched_clock.advance(0.1 + 1e-6)
     assert tracker.total == 200
-    time.sleep(0.1)
+    patched_clock.advance(0.2)
     assert tracker.total == 0