BUG: fix set_cosmology (#1005)

* BUG: fix set_cosmology

* TST: add test for getting cosmology after global is set

* TST: use fixture to ensure cosmology is reset

Author: mj-will

bilby/gw/cosmology.py

@@ -106,6 +106,8 @@ def set_cosmology(cosmology=None):
     """
     from ..core.utils.meta_data import global_meta_data
     cosmology = get_cosmology(cosmology)
+    global COSMOLOGY, DEFAULT_COSMOLOGY
+    DEFAULT_COSMOLOGY = cosmology
     COSMOLOGY[0] = cosmology
     if cosmology.name is not None:
         COSMOLOGY[1] = cosmology.name

test/gw/cosmology_test.py

@@ -1,8 +1,22 @@
 import unittest
 
-from astropy.cosmology import WMAP9, Planck15
+from astropy.cosmology import WMAP9
 from bilby.gw import cosmology
 import lal
+import pytest
+
+
+@pytest.fixture(autouse=True)
+def reset_cosmology():
+    # Reset cosmology before each test
+    cosmology.DEFAULT_COSMOLOGY = None
+    cosmology.COSMOLOGY = [None, str(None)]
+    try:
+        yield
+    finally:
+        # Reset cosmology after each test
+        cosmology.DEFAULT_COSMOLOGY = None
+        cosmology.COSMOLOGY = [None, str(None)]
 
 
 class TestSetCosmology(unittest.TestCase):
@@ -12,12 +26,11 @@ def setUp(self):
     def test_setting_cosmology_with_string(self):
         cosmology.set_cosmology("WMAP9")
         self.assertEqual(cosmology.COSMOLOGY[1], "WMAP9")
-        cosmology.set_cosmology("Planck15")
+        self.assertEqual(cosmology.DEFAULT_COSMOLOGY.name, "WMAP9")
 
     def test_setting_cosmology_with_astropy_object(self):
         cosmology.set_cosmology(WMAP9)
         self.assertEqual(cosmology.COSMOLOGY[1], "WMAP9")
-        cosmology.set_cosmology(Planck15)
 
     def test_setting_cosmology_with_default(self):
         cosmology.set_cosmology()
@@ -38,6 +51,15 @@ def test_setting_cosmology_with_w_cdm_dict(self):
         cosmology.set_cosmology(cosmo_dict)
         self.assertEqual(cosmology.COSMOLOGY[1][:4], "wCDM")
 
+    def test_repeated_setting_cosmology(self):
+        cosmology.set_cosmology("WMAP9")
+        self.assertEqual(cosmology.COSMOLOGY[1], "WMAP9")
+        self.assertEqual(cosmology.DEFAULT_COSMOLOGY.name, "WMAP9")
+        cosmology.set_cosmology("Planck18")
+        self.assertEqual(cosmology.COSMOLOGY[1], "Planck18")
+        self.assertEqual(cosmology.DEFAULT_COSMOLOGY.name, "Planck18")
+        cosmology.set_cosmology("Planck15")
+
 
 class TestGetCosmology(unittest.TestCase):
     def setUp(self):
@@ -49,6 +71,10 @@ def test_getting_cosmology_with_string(self):
     def test_getting_cosmology_with_default(self):
         self.assertEqual(cosmology.get_cosmology().name, "Planck15")
 
+    def test_getting_cosmology_non_standard_default(self):
+        cosmology.set_cosmology("WMAP9")
+        self.assertEqual(cosmology.get_cosmology().name, "WMAP9")
+
 
 class TestPlanck15LALCosmology(unittest.TestCase):