Refactor NaN handling test

nabu/preproc/tests/test_flatfield.py

@@ -270,11 +270,7 @@ class TestFlatField:
         self.check_normalized_radios(radios, config["expected_result"])
 
 
-    def test_twoflats_numerical_issue(self):
-        """
-        Same as above, but for the first radio: I==Dark and Flat==Dark
-        For this radio, nan is replaced with 1.
-        """
+    def _setup_numerical_issue(self):
         radios, flats, darks, config = self.get_test_elements("two_flats_with_radios_indices")
 
         # Retrieve the actual data for radios/darks/flats to use FlatField instead of FlatFieldDataUrl.
@@ -286,6 +282,27 @@ class TestFlatField:
             darks[dark_idx] = get_data(dark_url)
             darks[dark_idx][0, 0] = 99
         radios[:, 0, 0] = 99
+        return radios, flats, darks, config
+
+    def _check_numerical_issue(self, radios, expected_result, nan_value=None):
+        if nan_value is None:
+            assert np.alltrue(np.logical_not(np.isfinite(radios[:, 0, 0]))), "First pixel should be nan or inf"
+            radios[:, 0, 0] = radios[:, 1, 1]
+            self.check_normalized_radios(radios, expected_result)
+        else:
+            assert np.all(np.isfinite(radios)), "No inf/nan value should be there"
+            assert np.allclose(radios[:, 0, 0], nan_value, atol=1e-7), "Handled NaN should have nan_value=%f" % nan_value
+            radios[:, 0, 0] = radios[:, 1, 1]
+            self.check_normalized_radios(radios, expected_result)
+
+
+
+    def test_twoflats_numerical_issue(self):
+        """
+        Same as above, but for the first radio: I==Dark and Flat==Dark
+        For this radio, nan is replaced with 1.
+        """
+        radios, flats, darks, config = self._setup_numerical_issue()
         radios0 = radios.copy()
 
         # FlatField without NaN handling yields NaN and raises RuntimeWarning
@@ -295,9 +312,7 @@ class TestFlatField:
         )
         with pytest.warns(RuntimeWarning):
             FF_no_nan_handling.normalize_radios(radios)
-        assert np.alltrue(np.logical_not(np.isfinite(radios[:, 0, 0]))), "First pixel should be nan or inf"
-        radios[:, 0, 0] = radios[:, 1, 1]
-        self.check_normalized_radios(radios, config["expected_result"])
+        self._check_numerical_issue(radios, config["expected_result"], None)
 
         # FlatField with NaN handling
         nan_value = 50
@@ -308,10 +323,9 @@ class TestFlatField:
         )
         with pytest.warns(RuntimeWarning):
             FF_with_nan_handling.normalize_radios(radios)
-        assert np.all(np.isfinite(radios)), "No inf/nan value should be there"
-        assert np.allclose(radios[:, 0, 0], nan_value, atol=1e-7), "Handled NaN should have nan_value=%f" % nan_value
-        radios[:, 0, 0] = radios[:, 1, 1]
-        self.check_normalized_radios(radios, config["expected_result"])
+        self._check_numerical_issue(radios, config["expected_result"], nan_value)
+
+
 
 
 # This test should be closer to the ESRF standard setting.