afdo_tools: make bisection test-cases use tempdirs

These can be easily handed off to test-cases through the environment.

BUG=b:338238707
TEST=./run_python_tests.sh

Change-Id: I97d5aa6785f9e9c39056ea9c86621c00a5a3fab1
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/5507033
Reviewed-by: Ryan Beltran <ryanbeltran@chromium.org>
Tested-by: George Burgess <gbiv@chromium.org>
Commit-Queue: George Burgess <gbiv@chromium.org>

3 files changed, 28 insertions, 22 deletions

diff --git a/afdo_tools/bisection/afdo_prof_analysis_e2e_test.py b/afdo_tools/bisection/afdo_prof_analysis_e2e_test.py
index 4390807e..50fc2d27 100644
--- a/afdo_tools/bisection/afdo_prof_analysis_e2e_test.py
+++ b/afdo_tools/bisection/afdo_prof_analysis_e2e_test.py
@@ -15,6 +15,7 @@ import tempfile
 import unittest
 
 from afdo_tools.bisection import afdo_prof_analysis as analysis
+from llvm_tools import test_helpers
 
 
 class ObjectWithFields(object):
@@ -31,7 +32,7 @@ class ObjectWithFields(object):
             setattr(self, key, val)
 
 
-class AfdoProfAnalysisE2ETest(unittest.TestCase):
+class AfdoProfAnalysisE2ETest(test_helpers.TempDirTestCase):
     """Class for end-to-end testing of AFDO Profile Analysis"""
 
     # nothing significant about the values, just easier to remember even vs odd
@@ -59,6 +60,27 @@ class AfdoProfAnalysisE2ETest(unittest.TestCase):
         "bisect_results": {"ranges": [], "individuals": ["func_a"]},
     }
 
+    def setUp(self):
+        super().setUp()
+
+        # Test scripts depend on AFDO_TEST_DIR pointing to a directory to run
+        # in. Set that up for them.
+        self.tempdir = self.make_tempdir()
+
+        saved_value = None
+        tmpdir_env_var = "AFDO_TEST_DIR"
+        saved_value = os.environ.get(tmpdir_env_var)
+        os.environ[tmpdir_env_var] = str(self.tempdir)
+
+        def restore_environ():
+            if saved_value is None:
+                del os.environ[tmpdir_env_var]
+            else:
+                os.environ[tmpdir_env_var] = saved_value
+
+        self.addCleanup(restore_environ)
+
+
     def test_afdo_prof_analysis(self):
         # Individual issues take precedence by nature of our algos
         # so first, that should be caught
@@ -165,14 +187,8 @@ class AfdoProfAnalysisE2ETest(unittest.TestCase):
         }
 
         my_dir = os.path.dirname(os.path.abspath(__file__))
-        # using a static temp dir rather than a dynamic one because these files
-        # are shared between the bash scripts and this Python test, and the
-        # arguments to the bash scripts are fixed by afdo_prof_analysis.py so
-        # it would be difficult to communicate dynamically generated directory
-        # to bash scripts
-        scripts_tmp_dir = os.path.join(my_dir, "afdo_test_tmp")
+        scripts_tmp_dir = os.path.join(self.tempdir, "afdo_test_tmp")
         os.mkdir(scripts_tmp_dir)
-        self.addCleanup(shutil.rmtree, scripts_tmp_dir, ignore_errors=True)
 
         # files used in the bash scripts used as external deciders below
         # - count_file tracks the current number of calls to the script in
@@ -193,10 +209,9 @@ class AfdoProfAnalysisE2ETest(unittest.TestCase):
             num_calls = int(f.read())
         os.remove(count_file)  # reset counts for second run
         finished_state_file = os.path.join(
-            my_dir,
+            self.tempdir,
             f"afdo_analysis_state.json.completed.{date.today()}",
         )
-        self.addCleanup(os.remove, finished_state_file)
 
         # runs the same analysis but interrupted each iteration
         interrupt_decider = os.path.join(
@@ -266,16 +281,7 @@ class AfdoProfAnalysisE2ETest(unittest.TestCase):
         # FIXME: This test ideally shouldn't be writing to the directory of
         # this file.
         if state_file is None:
-            state_file = os.path.join(dir_path, "afdo_analysis_state.json")
-
-            def rm_state():
-                try:
-                    os.unlink(state_file)
-                except OSError:
-                    # Probably because the file DNE. That's fine.
-                    pass
-
-            self.addCleanup(rm_state)
+            state_file = os.path.join(self.tempdir, "afdo_analysis_state.json")
 
         actual = analysis.main_impl(
             ObjectWithFields(
diff --git a/afdo_tools/bisection/state_assumption_external.sh b/afdo_tools/bisection/state_assumption_external.sh
index ecc48af0..47168f11 100755
--- a/afdo_tools/bisection/state_assumption_external.sh
+++ b/afdo_tools/bisection/state_assumption_external.sh
@@ -12,7 +12,7 @@ SKIP_STATUS=125
 PROBLEM_STATUS=127
 
 my_dir="$(dirname "$(readlink -m "$0")")"
-tmp_dir="${my_dir}/afdo_test_tmp"
+tmp_dir="${AFDO_TEST_DIR:?}/afdo_test_tmp"
 count_file=${tmp_dir}/.count
 
 # keep count for purpose of filenames
diff --git a/afdo_tools/bisection/state_assumption_interrupt.sh b/afdo_tools/bisection/state_assumption_interrupt.sh
index 6eb5c2cb..c96b7597 100755
--- a/afdo_tools/bisection/state_assumption_interrupt.sh
+++ b/afdo_tools/bisection/state_assumption_interrupt.sh
@@ -9,7 +9,7 @@
 PROBLEM_STATUS=127
 
 my_dir="$(dirname "$(readlink -m "$0")")"
-tmp_dir="${my_dir}/afdo_test_tmp"
+tmp_dir="${AFDO_TEST_DIR:?}/afdo_test_tmp"
 
 count_file="${tmp_dir}/.count"
 if [[ -f "${count_file}" ]]; then