From b96510fbb36ac62df5866fe197541a504bd69ed2 Mon Sep 17 00:00:00 2001
From: Stefan Agner <stefan@agner.ch>
Date: Thu, 10 Oct 2024 00:16:27 +0200
Subject: [PATCH] Fix pytest

---
 supervisor/jobs/decorator.py  |  2 +-
 tests/test_supervisor.py      | 29 ++++++++++++++++++++++++-----
 tests/utils/test_decorator.py | 15 +++++++++++++++
 3 files changed, 40 insertions(+), 6 deletions(-)
 create mode 100644 tests/utils/test_decorator.py

diff --git a/supervisor/jobs/decorator.py b/supervisor/jobs/decorator.py
index 5f8934b2193..f4410ca5532 100644
--- a/supervisor/jobs/decorator.py
+++ b/supervisor/jobs/decorator.py
@@ -332,7 +332,7 @@ async def wrapper(
                     and self.cleanup
                 ):
                     self.sys_jobs.remove_job(job)
-
+        wrapper._job_decorator = self
         return wrapper
 
     @staticmethod
diff --git a/tests/test_supervisor.py b/tests/test_supervisor.py
index 0254de9c6f8..2d1028c2f7d 100644
--- a/tests/test_supervisor.py
+++ b/tests/test_supervisor.py
@@ -1,6 +1,6 @@
 """Test supervisor object."""
 
-from datetime import datetime
+from datetime import datetime, timedelta
 import errno
 from unittest.mock import AsyncMock, Mock, PropertyMock, patch
 
@@ -8,6 +8,7 @@
 from aiohttp.client_exceptions import ClientError
 from awesomeversion import AwesomeVersion
 import pytest
+from time_machine import travel
 
 from supervisor.const import UpdateChannel
 from supervisor.coresys import CoreSys
@@ -18,9 +19,11 @@
     SupervisorUpdateError,
 )
 from supervisor.host.apparmor import AppArmorControl
+from supervisor.jobs.decorator import Job
 from supervisor.resolution.const import ContextType, IssueType
 from supervisor.resolution.data import Issue
 from supervisor.supervisor import Supervisor
+from tests.utils.test_decorator import get_job_decorator
 
 
 @pytest.fixture(name="websession", scope="function")
@@ -58,21 +61,37 @@ async def test_connectivity_check(
     assert supervisor_unthrottled.connectivity is connectivity
 
 
-@pytest.mark.parametrize("side_effect,call_count", [(ClientError(), 3), (None, 1)])
+@pytest.mark.parametrize(
+    "side_effect,call_interval,throttled",
+    [
+        (None, timedelta(minutes=5), True),
+        (None, timedelta(minutes=15), False),
+        (ClientError(), timedelta(seconds=20), True),
+        (ClientError(), timedelta(seconds=40), False),
+    ],
+)
 async def test_connectivity_check_throttling(
     coresys: CoreSys,
     websession: AsyncMock,
     side_effect: Exception | None,
-    call_count: int,
+    call_interval: timedelta,
+    throttled: bool,
 ):
     """Test connectivity check throttled when checks succeed."""
     coresys.supervisor.connectivity = None
     websession.head.side_effect = side_effect
 
-    for _ in range(3):
+    job_decorator = get_job_decorator(Supervisor.check_connectivity)
+    job_decorator._last_call[None] = datetime.min
+
+    with (
+        travel(datetime.now(), tick=False) as traveller,
+    ):
+        await coresys.supervisor.check_connectivity()
+        traveller.shift(call_interval)
         await coresys.supervisor.check_connectivity()
 
-    assert websession.head.call_count == call_count
+    assert websession.head.call_count == (1 if throttled else 2)
 
 
 async def test_update_failed(coresys: CoreSys, capture_exception: Mock):
diff --git a/tests/utils/test_decorator.py b/tests/utils/test_decorator.py
new file mode 100644
index 00000000000..df00a29705a
--- /dev/null
+++ b/tests/utils/test_decorator.py
@@ -0,0 +1,15 @@
+"""Helper for decorators."""
+
+from supervisor.jobs.decorator import Job
+
+
+def get_job_decorator(func):
+    """Get Job object of decorated function."""
+    # Access the closure of the wrapper function
+    closure = func.__closure__
+    if closure:
+        for cell in closure:
+            obj = cell.cell_contents
+            if isinstance(obj, Job):
+                return obj
+    return None