From 25c86395655086fa7bdfd8d05278854c81feea3c Mon Sep 17 00:00:00 2001
From: MilesCranmer <miles.cranmer@gmail.com>
Date: Fri, 6 Dec 2024 14:08:55 +0000
Subject: [PATCH] test: make registry test more robust

---
 pysr/test/test_startup.py | 43 ++++++++++++++++++---------------------
 1 file changed, 20 insertions(+), 23 deletions(-)

diff --git a/pysr/test/test_startup.py b/pysr/test/test_startup.py
index 5e2ce350..4b2a450b 100644
--- a/pysr/test/test_startup.py
+++ b/pysr/test/test_startup.py
@@ -165,6 +165,16 @@ class TestRegistryHelper(unittest.TestCase):
 
     def setUp(self):
         self.old_value = os.environ.get(PREFERENCE_KEY, None)
+        self.recorded_env_vars = []
+        self.hits = 0
+
+        def failing_operation():
+            self.recorded_env_vars.append(os.environ[PREFERENCE_KEY])
+            self.hits += 1
+            # Just add some package I know will not exist and also not be in the dependency chain:
+            jl.Pkg.add(name="AirspeedVelocity", version="100.0.0")
+
+        self.failing_operation = failing_operation
 
     def tearDown(self):
         if self.old_value is not None:
@@ -182,16 +192,10 @@ def test_non_julia_errors_reraised(self):
 
     def test_julia_error_triggers_fallback(self):
         os.environ[PREFERENCE_KEY] = "conservative"
-        recorded_env_vars = []
-
-        def failing_operation():
-            recorded_env_vars.append(os.environ[PREFERENCE_KEY])
-            # Just add some package I know will not exist and also not be in the dependency chain:
-            jl.Pkg.add(name="AirspeedVelocity", version="100.0.0")
 
         with self.assertWarns(Warning) as warn_context:
             with self.assertRaises(Exception) as error_context:
-                try_with_registry_fallback(failing_operation)
+                try_with_registry_fallback(self.failing_operation)
 
         self.assertIn(
             "Unsatisfiable requirements detected", str(error_context.exception)
@@ -202,7 +206,8 @@ def failing_operation():
         )
 
         # Verify both modes are tried in order
-        self.assertEqual(recorded_env_vars, ["conservative", "eager"])
+        self.assertEqual(self.recorded_env_vars, ["conservative", "eager"])
+        self.assertEqual(self.hits, 2)
 
         # Verify environment is restored
         self.assertEqual(os.environ[PREFERENCE_KEY], "conservative")
@@ -210,22 +215,14 @@ def failing_operation():
     def test_eager_mode_fails_directly(self):
         os.environ[PREFERENCE_KEY] = "eager"
 
-        recorded_env_vars = []
-        hits = 0
-
-        def failing_operation():
-            recorded_env_vars.append(os.environ[PREFERENCE_KEY])
-            nonlocal hits
-            hits += 1
-            raise Exception(
-                "JuliaError\nUnsatisfiable requirements detected for package Test [8dfed614-e22c-5e08-85e1-65c5234f0b40]"
-            )
-
         with self.assertRaises(Exception) as context:
-            try_with_registry_fallback(failing_operation)
-        self.assertIn("JuliaError", str(context.exception))
-        self.assertEqual(recorded_env_vars, ["eager"])  # Should only try eager mode
-        self.assertEqual(hits, 1)
+            try_with_registry_fallback(self.failing_operation)
+
+        self.assertIn("Unsatisfiable requirements detected", str(context.exception))
+        self.assertEqual(
+            self.recorded_env_vars, ["eager"]
+        )  # Should only try eager mode
+        self.assertEqual(self.hits, 1)
 
 
 def runtests(just_tests=False):