Fix race condition in test_psf.py

LSSTDESC · May 10, 2024 · 064a1a2 · 064a1a2
1 parent 993512f
commit 064a1a2
Showing 1 changed file with 58 additions and 69 deletions.
diff --git a/tests/test_psf.py b/tests/test_psf.py
@@ -1,14 +1,13 @@
 """
 Unit tests for PSF-related code.
 """
-import os
 import numpy as np
 from pathlib import Path
-import glob
 import unittest
 import time
 import imsim
 import galsim
+import tempfile
 
 DATA_DIR = Path(__file__).parent / 'data'
 
@@ -20,16 +19,8 @@ class PsfTestCase(unittest.TestCase):
     TestCase class for PSF-related functions.
     """
     def setUp(self):
-        self.test_dir = 'psf_tests_dir'
         instcat = DATA_DIR / 'tiny_instcat.txt'
         self.opsim_data = imsim.OpsimDataLoader(str(instcat))
-        if not os.path.exists(self.test_dir):
-            os.makedirs(self.test_dir)
-
-    def tearDown(self):
-        for item in glob.glob(os.path.join(self.test_dir, '*')):
-            os.remove(item)
-        os.rmdir(self.test_dir)
 
     def test_save_and_load_psf(self):
         """
@@ -85,77 +76,75 @@ def test_save_and_load_psf(self):
         }
         galsim.config.ProcessInput(config)
         config['wcs'] = galsim.config.BuildWCS(config['image'], 'wcs', config)
-        for psf_name in ("DoubleGaussian", "Kolmogorov", "Atmospheric"):
-            psf = galsim.config.BuildGSObject(config, psf_name)
-            psf_file = os.path.join(self.test_dir, '{}.pkl'.format(psf_name))
-            imsim.save_psf(psf, psf_file)
-            psf_retrieved = imsim.load_psf(psf_file)
-            self.assertEqual(psf, psf_retrieved)
+        with tempfile.NamedTemporaryFile() as psf_fp:
+            for psf_name in ("DoubleGaussian", "Kolmogorov", "Atmospheric"):
+                psf = galsim.config.BuildGSObject(config, psf_name)
+                imsim.save_psf(psf, psf_fp.name)
+                psf_retrieved = imsim.load_psf(psf_fp.name)
+                self.assertEqual(psf, psf_retrieved)
 
     def test_atm_psf_save_file(self):
         """Test using save_file in  AtmosphericPSF
         """
-        psf_file = os.path.join(self.test_dir, 'save_atm_psf.pkl')
-        config = {
-            'psf': {
-                'type': 'AtmosphericPSF'
-            },
-            'input': {
-                'atm_psf': {
-                    'airmass': self.opsim_data['airmass'],
-                    'rawSeeing': self.opsim_data['rawSeeing'],
-                    'band':  self.opsim_data['band'],
-                    'screen_scale': 6.4,
-                    'boresight': {
-                        'type': 'RADec',
-                        'ra': { 'type': 'Degrees', 'theta': self.opsim_data['rightascension'], },
-                        'dec': { 'type': 'Degrees', 'theta': self.opsim_data['declination'], }
-                    },
-                    'save_file': psf_file
-                }
-            },
-            'image_pos': galsim.PositionD(0,0),  # This would get set appropriately during
-                                                 # normal config processing.
-            'image' : {
-                'random_seed': 1234,
-                'wcs': {
-                    'type' : 'Tan',
-                    'dudx' : 0.2,
-                    'dudy' : 0.,
-                    'dvdx' : 0.,
-                    'dvdy' : 0.2,
-                    'ra' : '@input.atm_psf.boresight.ra',
-                    'dec' : '@input.atm_psf.boresight.dec',
+        with tempfile.TemporaryDirectory() as tmpdir:
+            psf_file = Path(tmpdir) / 'atm_psf.pkl'
+            config = {
+                'psf': {
+                    'type': 'AtmosphericPSF'
+                },
+                'input': {
+                    'atm_psf': {
+                        'airmass': self.opsim_data['airmass'],
+                        'rawSeeing': self.opsim_data['rawSeeing'],
+                        'band':  self.opsim_data['band'],
+                        'screen_scale': 6.4,
+                        'boresight': {
+                            'type': 'RADec',
+                            'ra': { 'type': 'Degrees', 'theta': self.opsim_data['rightascension'], },
+                            'dec': { 'type': 'Degrees', 'theta': self.opsim_data['declination'], }
+                        },
+                        'save_file': psf_file
+                    }
+                },
+                'image_pos': galsim.PositionD(0,0),  # This would get set appropriately during
+                                                    # normal config processing.
+                'image' : {
+                    'random_seed': 1234,
+                    'wcs': {
+                        'type' : 'Tan',
+                        'dudx' : 0.2,
+                        'dudy' : 0.,
+                        'dvdx' : 0.,
+                        'dvdy' : 0.2,
+                        'ra' : '@input.atm_psf.boresight.ra',
+                        'dec' : '@input.atm_psf.boresight.dec',
+                    }
                 }
             }
-        }
-
-        if os.path.isfile(psf_file):
-            os.remove(psf_file)
 
-        config['wcs'] = galsim.config.BuildWCS(config['image'], 'wcs', config)
-        config1 = galsim.config.CopyConfig(config)
-        config2 = galsim.config.CopyConfig(config)
+            config['wcs'] = galsim.config.BuildWCS(config['image'], 'wcs', config)
+            config1 = galsim.config.CopyConfig(config)
+            config2 = galsim.config.CopyConfig(config)
 
-        # The first time, it will build the psf from scratch and save the screens.
-        t0 = time.time()
-        galsim.config.ProcessInput(config1)
-        t1 = time.time()
+            # The first time, it will build the psf from scratch and save the screens.
+            t0 = time.time()
+            galsim.config.ProcessInput(config1)
+            t1 = time.time()
 
-        assert os.path.isfile(psf_file)
+            assert psf_file.is_file()
 
-        # The second time, it will be faster, since it loads the screens from the file.
-        t2 = time.time()
-        galsim.config.ProcessInput(config2)
-        t3 = time.time()
+            # The second time, it will be faster, since it loads the screens from the file.
+            t2 = time.time()
+            galsim.config.ProcessInput(config2)
+            t3 = time.time()
 
-        print('Times = ',t1-t0,t3-t2)
-        assert t1-t0 > t3-t2
+            print('Times = ',t1-t0,t3-t2)
+            assert t1-t0 > t3-t2
 
-        # Both input objects will make the same PSF at the same location:
-        psf1 = galsim.config.BuildGSObject(config1, 'psf')[0]
-        psf2 = galsim.config.BuildGSObject(config2, 'psf')[0]
-        assert psf1 == psf2
+            # Both input objects will make the same PSF at the same location:
+            psf1 = galsim.config.BuildGSObject(config1, 'psf')[0]
+            psf2 = galsim.config.BuildGSObject(config2, 'psf')[0]
+            assert psf1 == psf2
 
 
     def test_atm_psf_config(self):