Просмотр исходного кода

ref: fix permissions and tempdir in test_models (#52725)

<!-- Describe your PR here. -->
anthony sottile 1 год назад
Родитель
Сommit
8cdd25ffae
1 измененных файлов с 22 добавлено и 29 удалено
  1. 22 29
      tests/sentry/backup/test_models.py

+ 22 - 29
tests/sentry/backup/test_models.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import atexit
 import tempfile
 from pathlib import Path
 
@@ -21,41 +20,35 @@ class ModelBackupTests(TransactionTestCase):
     comparators."""
 
     def setUp(self):
-        # Create a temporary directory for JSON exports.
-        self.tmp_dir = tempfile.TemporaryDirectory()
-        atexit.register(self.tmp_dir.cleanup)
-
-        # Reset the Django database.
-        call_command("flush", verbosity=0, interactive=False)
-
-        # Generate temporary filenames for the expected and actual JSON files.
-        self.tmp_expect = Path(self.tmp_dir.name) / f"{self._testMethodName}.expect.json"
-        self.tmp_actual = Path(self.tmp_dir.name) / f"{self._testMethodName}.actual.json"
-
-    def tearDown(self):
-        self.tmp_dir.cleanup()
+        with in_test_psql_role_override("postgres"):
+            # Reset the Django database.
+            call_command("flush", verbosity=0, interactive=False)
 
     def import_export_then_validate(self):
         """Test helper that validates that data imported from a temporary `.json` file correctly
         matches the actual outputted export data."""
 
-        # Export the current state of the database into the "expected" temporary file, then parse it
-        # into a JSON object for comparison.
-        expect = tmp_export_to_file(self.tmp_expect)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            tmp_expect = Path(tmpdir).joinpath(f"{self._testMethodName}.expect.json")
+            tmp_actual = Path(tmpdir).joinpath(f"{self._testMethodName}.actual.json")
 
-        # Reset the Django database.
-        call_command("flush", verbosity=0, interactive=False)
+            # Export the current state of the database into the "expected" temporary file, then parse it
+            # into a JSON object for comparison.
+            expect = tmp_export_to_file(tmp_expect)
 
-        # Write the contents of the "expected" JSON file into the now clean database.
-        with in_test_psql_role_override("postgres"):
-            rv = CliRunner().invoke(import_, [str(self.tmp_expect)])
-            assert rv.exit_code == 0, rv.output
-
-        # Validate that the "expected" and "actual" JSON matches.
-        actual = tmp_export_to_file(self.tmp_actual)
-        res = validate(expect, actual)
-        if res.findings:
-            raise ValidationError(res)
+            # Write the contents of the "expected" JSON file into the now clean database.
+            with in_test_psql_role_override("postgres"):
+                # Reset the Django database.
+                call_command("flush", verbosity=0, interactive=False)
+
+                rv = CliRunner().invoke(import_, [str(tmp_expect)])
+                assert rv.exit_code == 0, rv.output
+
+            # Validate that the "expected" and "actual" JSON matches.
+            actual = tmp_export_to_file(tmp_actual)
+            res = validate(expect, actual)
+            if res.findings:
+                raise ValidationError(res)
 
     def test_organization(self):
         user = self.create_user()