add tests for the new gain replacement functions

tests/test_jungfraulib.py

+import numpy as np
 import pytest
 from extra_data import RunDirectory
 
-from cal_tools.jungfraulib import JungfrauCtrl, sort_runs_by_gain
+from cal_tools.enums import BadPixels
+from cal_tools.jungfraulib import (
+    JungfrauCtrl,
+    replace_pixels_gain_values,
+    replace_pixels_wrong_gain_badpixel,
+    sort_runs_by_gain,
+)
 
 # TODO: replace with mocked RAW data as in tests/test_agipdlib.py
 JF = JungfrauCtrl(
@@ -64,3 +71,64 @@ def test_sort_runs_by_gain(original_runs, sorted_runs):
         runs=original_runs,
         ctrl_src="FXE_XAD_JF1M/DET/CONTROL")
     assert validated_runs == sorted_runs
+
+
+GAIN = np.zeros((1, 512, 1024), dtype=np.uint8)
+NEW_GAIN = GAIN.copy()  # all pixels in high gain.
+GAIN[:, 0:255, 896:1024] = 3  # subsection of a chip is stuck in gain bit = 3
+BP = np.zeros((512, 1024, 3), dtype=np.uint32)
+for i in range(3):
+    BP[0:255, 896:1024, i] = BadPixels.WRONG_GAIN_VALUE.value
+
+MEM_CELLS = 1
+KARABO_DA = "JNGFR04"
+REPLACE_WRONG_GAIN_VALUE = 0
+WRONG_GAIN_PIXELS = [4, 0, 255, 896, 1024]
+
+
+def test_replace_pixels_wrong_gain_badpixel():
+    assert np.allclose(
+        NEW_GAIN,
+        replace_pixels_wrong_gain_badpixel(
+            gain_arr=GAIN,
+            new_gain_value=REPLACE_WRONG_GAIN_VALUE,
+            badpixel_map=BP,
+            memory_cells=MEM_CELLS,
+        )
+    )
+
+
+def test_replace_pixels_gain_values():
+    assert np.allclose(
+        NEW_GAIN,
+        replace_pixels_gain_values(
+            gain_arr=GAIN,
+            new_gain_value=REPLACE_WRONG_GAIN_VALUE,
+            wrong_gain_pixels=WRONG_GAIN_PIXELS,
+            karabo_da=KARABO_DA,
+            badpixel_map=BP,
+            memory_cells=MEM_CELLS,
+        )
+    )
+    assert np.allclose(
+        NEW_GAIN,
+        replace_pixels_gain_values(
+            gain_arr=GAIN,
+            new_gain_value=REPLACE_WRONG_GAIN_VALUE,
+            wrong_gain_pixels=WRONG_GAIN_PIXELS,
+            karabo_da=KARABO_DA,
+            badpixel_map=BP,
+            memory_cells=MEM_CELLS,
+        )
+    )
+    assert np.allclose(
+        GAIN,  # No replacement for non-selected module.
+        replace_pixels_gain_values(
+            gain_arr=GAIN,
+            new_gain_value=REPLACE_WRONG_GAIN_VALUE,
+            wrong_gain_pixels=WRONG_GAIN_PIXELS,
+            karabo_da="JNGFR01",
+            badpixel_map=BP,
+            memory_cells=MEM_CELLS,
+        )
+    )