From 95cb43f55d127df02532af493c6c46024514b104 Mon Sep 17 00:00:00 2001
From: Yisu Peng <shawn_pys@yahoo.com>
Date: Fri, 18 Aug 2023 03:23:41 -0400
Subject: [PATCH] Remove temporary directories after test done (#6236)

* remove temporary isolated venv after isolated notebook/packaging test done

* remove temporary venv after cloned env test done

* remove temporary directories for figures after heatmap testings

Fixes #6036
---
 cirq-core/cirq/vis/heatmap_test.py            | 2 ++
 dev_tools/cloned_env_test.py                  | 2 ++
 dev_tools/notebooks/isolated_notebook_test.py | 2 ++
 dev_tools/packaging/isolated_packages_test.py | 2 ++
 4 files changed, 8 insertions(+)

diff --git a/cirq-core/cirq/vis/heatmap_test.py b/cirq-core/cirq/vis/heatmap_test.py
index 252afd0edf4..20a0b4eea39 100644
--- a/cirq-core/cirq/vis/heatmap_test.py
+++ b/cirq-core/cirq/vis/heatmap_test.py
@@ -14,6 +14,7 @@
 """Tests for Heatmap."""
 
 import pathlib
+import shutil
 import string
 from tempfile import mkdtemp
 
@@ -309,6 +310,7 @@ def test_colorbar(ax, position, size, pad):
 
     plt.close(fig1)
     plt.close(fig2)
+    shutil.rmtree(tmp_dir)
 
 
 @pytest.mark.usefixtures('closefigures')
diff --git a/dev_tools/cloned_env_test.py b/dev_tools/cloned_env_test.py
index 883737e7cb0..5284d05ee16 100644
--- a/dev_tools/cloned_env_test.py
+++ b/dev_tools/cloned_env_test.py
@@ -15,6 +15,7 @@
 """Tests the cloned_env fixture in conftest.py"""
 import json
 import os
+import shutil
 import subprocess
 from unittest import mock
 
@@ -39,3 +40,4 @@ def test_isolated_env_cloning(cloned_env, param):
     packages = json.loads(result.stdout)
     assert {"name": "flynt", "version": "0.64"} in packages
     assert {"astor", "flynt", "pip", "setuptools", "wheel"} == set(p['name'] for p in packages)
+    shutil.rmtree(env)
diff --git a/dev_tools/notebooks/isolated_notebook_test.py b/dev_tools/notebooks/isolated_notebook_test.py
index 03e03a9e15b..a1f1e7f9912 100644
--- a/dev_tools/notebooks/isolated_notebook_test.py
+++ b/dev_tools/notebooks/isolated_notebook_test.py
@@ -27,6 +27,7 @@
 import os
 import re
 import subprocess
+import shutil
 import warnings
 from typing import Set, List
 
@@ -188,6 +189,7 @@ def _rewrite_and_run_notebook(notebook_path, cloned_env):
             f"dev_tools/notebooks/isolated_notebook_test.py."
         )
     os.remove(rewritten_notebook_path)
+    shutil.rmtree(notebook_env)
 
 
 @pytest.mark.slow
diff --git a/dev_tools/packaging/isolated_packages_test.py b/dev_tools/packaging/isolated_packages_test.py
index c3c92e7a61b..3f03d2408fe 100644
--- a/dev_tools/packaging/isolated_packages_test.py
+++ b/dev_tools/packaging/isolated_packages_test.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+import shutil
 import subprocess
 from unittest import mock
 
@@ -48,3 +49,4 @@ def test_isolated_packages(cloned_env, module):
         check=False,
     )
     assert result.returncode == 0, f"Failed isolated tests for {module.name}:\n{result.stdout}"
+    shutil.rmtree(env)