From 1254c319da5b2f66cf3fc06182fe4dea37978b81 Mon Sep 17 00:00:00 2001
From: Mike Degatano <michael.degatano@gmail.com>
Date: Thu, 10 Oct 2024 18:35:04 +0000
Subject: [PATCH] Reset throttle in test and refactor helper

---
 tests/common.py                  | 17 +++++++++++++++++
 tests/jobs/test_job_decorator.py |  4 ++++
 tests/test_supervisor.py         |  6 ++----
 tests/utils/test_decorator.py    | 15 ---------------
 4 files changed, 23 insertions(+), 19 deletions(-)
 delete mode 100644 tests/utils/test_decorator.py

diff --git a/tests/common.py b/tests/common.py
index ad489a566a2..d7b507e9757 100644
--- a/tests/common.py
+++ b/tests/common.py
@@ -1,12 +1,15 @@
 """Common test functions."""
 
+from datetime import datetime
 from importlib import import_module
+from inspect import getclosurevars
 import json
 from pathlib import Path
 from typing import Any
 
 from dbus_fast.aio.message_bus import MessageBus
 
+from supervisor.jobs.decorator import Job
 from supervisor.resolution.validate import get_valid_modules
 from supervisor.utils.yaml import read_yaml_file
 
@@ -82,3 +85,17 @@ async def mock_dbus_services(
                 services[module] = service_module.setup(to_mock[module]).export(bus)
 
     return services
+
+
+def get_job_decorator(func) -> Job:
+    """Get Job object of decorated function."""
+    # Access the closure of the wrapper function
+    job = getclosurevars(func).nonlocals["self"]
+    if not isinstance(job, Job):
+        raise TypeError(f"{func.__qualname__} is not a Job")
+    return job
+
+
+def reset_last_call(func, group: str | None = None) -> None:
+    """Reset last call for a function using the Job decorator."""
+    get_job_decorator(func).set_last_call(datetime.min, group)
diff --git a/tests/jobs/test_job_decorator.py b/tests/jobs/test_job_decorator.py
index 426973d7c75..8d618d94805 100644
--- a/tests/jobs/test_job_decorator.py
+++ b/tests/jobs/test_job_decorator.py
@@ -26,8 +26,11 @@
 from supervisor.os.manager import OSManager
 from supervisor.plugins.audio import PluginAudio
 from supervisor.resolution.const import UnhealthyReason
+from supervisor.supervisor import Supervisor
 from supervisor.utils.dt import utcnow
 
+from tests.common import reset_last_call
+
 
 async def test_healthy(coresys: CoreSys, caplog: pytest.LogCaptureFixture):
     """Test the healty decorator."""
@@ -73,6 +76,7 @@ async def test_internet(
 ):
     """Test the internet decorator."""
     coresys.core.state = CoreState.RUNNING
+    reset_last_call(Supervisor.check_connectivity)
 
     class TestClass:
         """Test class."""
diff --git a/tests/test_supervisor.py b/tests/test_supervisor.py
index f2fe4917b89..9e093fc0fa7 100644
--- a/tests/test_supervisor.py
+++ b/tests/test_supervisor.py
@@ -23,7 +23,7 @@
 from supervisor.resolution.data import Issue
 from supervisor.supervisor import Supervisor
 
-from tests.utils.test_decorator import get_job_decorator
+from tests.common import reset_last_call
 
 
 @pytest.fixture(name="websession", scope="function")
@@ -81,9 +81,7 @@ async def test_connectivity_check_throttling(
     coresys.supervisor.connectivity = None
     websession.head.side_effect = side_effect
 
-    job_decorator = get_job_decorator(Supervisor.check_connectivity)
-    job_decorator._last_call.clear()  # pylint: disable=W0212
-
+    reset_last_call(Supervisor.check_connectivity)
     with (
         travel(datetime.now(), tick=False) as traveller,
     ):
diff --git a/tests/utils/test_decorator.py b/tests/utils/test_decorator.py
deleted file mode 100644
index df00a29705a..00000000000
--- a/tests/utils/test_decorator.py
+++ /dev/null
@@ -1,15 +0,0 @@
-"""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