athena tests: patch instead of juggling function (#30664)

patch instead of juggling function

selfdrive/athena/tests/test_athenad_ping.py

@@ -3,8 +3,8 @@ import subprocess
 import threading
 import time
 import unittest
-from typing import Callable, cast, Optional
-from unittest.mock import MagicMock
+from typing import cast, Optional
+from unittest import mock
 
 from openpilot.common.params import Params
 from openpilot.common.timeout import Timeout
@@ -27,8 +27,6 @@ class TestAthenadPing(unittest.TestCase):
   athenad: threading.Thread
   exit_event: threading.Event
 
-  _create_connection: Callable
-
   def _get_ping_time(self) -> Optional[str]:
     return cast(Optional[str], self.params.get("LastAthenaPingTime", encoding="utf-8"))
 
@@ -38,15 +36,9 @@ class TestAthenadPing(unittest.TestCase):
   def _received_ping(self) -> bool:
     return self._get_ping_time() is not None
 
-  @classmethod
-  def setUpClass(cls) -> None:
-    cls._create_connection = athenad.create_connection
-    athenad.create_connection = MagicMock(wraps=cls._create_connection)
-
   @classmethod
   def tearDownClass(cls) -> None:
     wifi_radio(True)
-    athenad.create_connection = cls._create_connection
 
   def setUp(self) -> None:
     self.params = Params()
@@ -58,19 +50,18 @@ class TestAthenadPing(unittest.TestCase):
     self.exit_event = threading.Event()
     self.athenad = threading.Thread(target=athenad.main, args=(self.exit_event,))
 
-    athenad.create_connection.reset_mock()
-
   def tearDown(self) -> None:
     if self.athenad.is_alive():
       self.exit_event.set()
       self.athenad.join()
 
-  def assertTimeout(self, reconnect_time: float) -> None:
+  @mock.patch('openpilot.selfdrive.athena.athenad.create_connection', autospec=True)
+  def assertTimeout(self, reconnect_time: float, mock_create_connection: mock.MagicMock) -> None:
     self.athenad.start()
 
     time.sleep(1)
-    athenad.create_connection.assert_called_once()
-    athenad.create_connection.reset_mock()
+    mock_create_connection.assert_called_once()
+    mock_create_connection.reset_mock()
 
     # check normal behaviour
     with self.subTest("Wi-Fi: receives ping"), Timeout(70, "no ping received"):
@@ -78,7 +69,7 @@ class TestAthenadPing(unittest.TestCase):
         time.sleep(0.1)
       print("ping received")
 
-    athenad.create_connection.assert_not_called()
+    mock_create_connection.assert_not_called()
 
     # websocket should attempt reconnect after short time
     with self.subTest("LTE: attempt reconnect"):
@@ -86,7 +77,7 @@ class TestAthenadPing(unittest.TestCase):
       print("waiting for reconnect attempt")
       start_time = time.monotonic()
       with Timeout(reconnect_time, "no reconnect attempt"):
-        while not athenad.create_connection.called:
+        while not mock_create_connection.called:
           time.sleep(0.1)
         print(f"reconnect attempt after {time.monotonic() - start_time:.2f}s")