From 1cc555c37bb1e4a10780a1946404a1f14f492d4c Mon Sep 17 00:00:00 2001
From: Justin Newberry <justin@comma.ai>
Date: Thu, 31 Aug 2023 14:00:31 -0700
Subject: [PATCH] Tests: expand temporary directory helpers for more
 directories (#29735)

* create helpers for testing

* document it

* fix other tests
old-commit-hash: 405d09714c15956a6d220ff963d07b89c8e790ca
---
 selfdrive/test/helpers.py       | 42 +++++++++++++++++++++++++++------
 tools/lib/tests/test_caching.py |  4 ++--
 2 files changed, 37 insertions(+), 9 deletions(-)

diff --git a/selfdrive/test/helpers.py b/selfdrive/test/helpers.py
index 3f4e4e1661..e376a86706 100644
--- a/selfdrive/test/helpers.py
+++ b/selfdrive/test/helpers.py
@@ -1,6 +1,7 @@
 import os
 import time
 import tempfile
+from typing import List, Union
 
 from unittest import mock
 from functools import wraps
@@ -72,17 +73,44 @@ def with_processes(processes, init_time=0, ignore_stopped=None):
   return wrapper
 
 
-def temporary_mock_dir(mock_path):
+def temporary_mock_dir(mock_paths_in: Union[List[str], str], kwarg: Union[str, None] = None, generator = tempfile.TemporaryDirectory):
+  """
+  mock_paths_in: string or string list representing the full path of the variable you want to mock.
+  kwarg: str or None representing the kwarg that gets passed into the test function, in case the test needs access to the temporary directory.
+  generator: a context to use to generate the temporary directory
+  """
   def wrapper(func):
     @wraps(func)
     def wrap(*args, **kwargs):
-      with tempfile.TemporaryDirectory() as temp_dir:
-        cache_dir_patch = mock.patch(mock_path, temp_dir)
-        cache_dir_patch.start()
-        func(*args, **kwargs, temp_dir=temp_dir)
-        cache_dir_patch.stop()
+      mock_paths = mock_paths_in if isinstance(mock_paths_in, list) else [mock_paths_in]
+      with generator() as temp_dir:
+        mocks = []
+        for mock_path in mock_paths:
+          mocks.append(mock.patch(mock_path, str(temp_dir)))
+        [mock.start() for mock in mocks]
+        try:
+          if kwarg is not None:
+            kwargs[kwarg] = temp_dir
+          func(*args, **kwargs)
+        finally:
+          [mock.stop() for mock in mocks]
+
     return wrap
   return wrapper
 
+def string_context(context):
+  class StringContext:
+    def __enter__(self):
+      return context
+
+    def __exit__(self, *args):
+      pass
+
+  return StringContext
+
+temporary_dir = temporary_mock_dir([], "temp_dir")
 temporary_cache_dir = temporary_mock_dir("openpilot.tools.lib.url_file.CACHE_DIR")
-temporary_swaglog_dir = temporary_mock_dir("openpilot.system.swaglog.SWAGLOG_DIR")
\ No newline at end of file
+temporary_swaglog_dir = temporary_mock_dir("openpilot.system.swaglog.SWAGLOG_DIR", "temp_dir")
+temporary_laikad_downloads_dir = temporary_mock_dir("openpilot.selfdrive.locationd.laikad.DOWNLOADS_CACHE_FOLDER")
+temporary_swaglog_ipc = temporary_mock_dir(["openpilot.system.swaglog.SWAGLOG_IPC", "system.logmessaged.SWAGLOG_IPC"],
+                                           generator=string_context("/tmp/test_swaglog_ipc"))
\ No newline at end of file
diff --git a/tools/lib/tests/test_caching.py b/tools/lib/tests/test_caching.py
index bd93cd93ec..7702b52e3e 100644
--- a/tools/lib/tests/test_caching.py
+++ b/tools/lib/tests/test_caching.py
@@ -32,7 +32,7 @@ class TestFileDownload(unittest.TestCase):
     self.assertEqual(response_cached, response_downloaded)
 
   @temporary_cache_dir
-  def test_small_file(self, temp_dir):
+  def test_small_file(self):
     # Make sure we don't force cache
     os.environ["FILEREADER_CACHE"] = "0"
     small_file_url = "https://raw.githubusercontent.com/commaai/openpilot/master/docs/SAFETY.md"
@@ -53,7 +53,7 @@ class TestFileDownload(unittest.TestCase):
       self.compare_loads(small_file_url, 100 * i, 100)
 
   @temporary_cache_dir
-  def test_large_file(self, temp_dir):
+  def test_large_file(self):
     large_file_url = "https://commadataci.blob.core.windows.net/openpilotci/0375fdf7b1ce594d/2019-06-13--08-32-25/3/qlog.bz2"
     #  Load the end 100 bytes of both files
     file_large = URLFile(large_file_url)