""" interpolator.py 多进程重构的行为测试（mock-based，无 osgeo 依赖） 验证： 1. 静态结构：模块级函数集合、签名、新参数、_worker_dataset 全局 2. 行为逻辑：_process_one_block 在 mock dataset 上的零像素识别 + 插值 3. 行为逻辑：_interpolate_block_worker 通过模块全局 _worker_dataset 工作 4. 行为逻辑：interpolate_zero_pixels_batch 的串行路径（不依赖 osgeo） 5. 向后兼容：现有 8 个 caller 参数保留默认值不变 如果本机有 osgeo，可补一个真实数据集的 smoke test。 """ import os import sys import ast import types import unittest from unittest.mock import MagicMock, patch # Ensure the project src is on path so we can import the module under test THIS_DIR = os.path.dirname(os.path.abspath(__file__)) PROJECT_ROOT = os.path.abspath(os.path.join(THIS_DIR, "..")) if PROJECT_ROOT not in sys.path: sys.path.insert(0, PROJECT_ROOT) INTERP_PATH = os.path.join( PROJECT_ROOT, "src", "core", "algorithms", "interpolation", "interpolator.py" ) # ============================================================================= # Static structure tests # ============================================================================= class TestInterpolatorStructure(unittest.TestCase): def setUp(self): with open(INTERP_PATH, "r", encoding="utf-8") as f: self.src = f.read() self.tree = ast.parse(self.src) def test_module_level_functions(self): mod_funcs = {n.name for n in self.tree.body if isinstance(n, ast.FunctionDef)} expected = { "interpolate_pixels", "_interpolate_single_band", "_normalize_interpolation_method", "_read_water_mask_to_array", "_init_worker", "_interpolate_block_worker", "_process_one_block", "interpolate_zero_pixels_batch", } self.assertEqual(mod_funcs, expected) self.assertNotIn("_process_block_with_buffer", mod_funcs) def test_worker_dataset_module_global(self): mod_globals = set() for n in self.tree.body: if isinstance(n, ast.Assign) and isinstance(n.targets[0], ast.Name): mod_globals.add(n.targets[0].id) elif isinstance(n, ast.AnnAssign) and isinstance(n.target, ast.Name): mod_globals.add(n.target.id) self.assertIn("_worker_dataset", mod_globals) def test_interpolate_zero_pixels_batch_signature(self): for n in self.tree.body: if isinstance(n, ast.FunctionDef) and n.name == "interpolate_zero_pixels_batch": args = [a.arg for a in n.args.args] # Backward compat: all 8 existing params + 2 new self.assertEqual( args, [ "img_path", "interpolation_method", "output_path", "water_mask", "deglint_dir", "callback_progress", "block_size", "halo_size", "n_workers", "use_multiprocessing", ], ) defaults = [getattr(d, "value", None) for d in n.args.defaults] self.assertEqual( defaults, ["nearest", None, None, None, None, 1024, 64, None, True], ) return self.fail("interpolate_zero_pixels_batch not found") def test_init_worker_signature(self): for n in self.tree.body: if isinstance(n, ast.FunctionDef) and n.name == "_init_worker": self.assertEqual([a.arg for a in n.args.args], ["img_path"]) return self.fail("_init_worker not found") def test_worker_function_signature(self): for n in self.tree.body: if isinstance(n, ast.FunctionDef) and n.name == "_interpolate_block_worker": self.assertEqual([a.arg for a in n.args.args], ["task"]) return self.fail("_interpolate_block_worker not found") def test_process_one_block_signature(self): for n in self.tree.body: if isinstance(n, ast.FunctionDef) and n.name == "_process_one_block": args = [a.arg for a in n.args.args] self.assertEqual( args, [ "dataset", "x0", "y0", "ey0", "ex0", "ey1", "ex1", "row_offset", "col_offset", "inner_h", "inner_w", "mask_segment_ext", "method", ], ) return self.fail("_process_one_block not found") def test_uses_process_pool_executor(self): self.assertIn("ProcessPoolExecutor", self.src) self.assertIn("ProcessPoolExecutor(", self.src) self.assertIn("initializer=_init_worker", self.src) self.assertIn("initargs=(img_path,)", self.src) self.assertIn("GDAL_NUM_THREADS", self.src) def test_dispatch_logic_present(self): # Both serial and parallel paths should be present self.assertIn("if effective_workers <= 1:", self.src) self.assertIn("with ProcessPoolExecutor", self.src) def test_serial_path_uses_process_one_block(self): # Serial branch should call _process_one_block directly (no pickle overhead) self.assertIn( "_process_one_block(\n dataset, *task\n )", self.src, ) def test_worker_path_uses_process_one_block(self): # Worker function should also call _process_one_block (shared core) self.assertIn("_process_one_block(", self.src) # The worker should read from _worker_dataset, not receive dataset in task self.assertIn("_worker_dataset", self.src) # ============================================================================= # Mocked behavior tests (no real osgeo/scipy needed) # ============================================================================= def _make_mock_band(read_data): band = MagicMock() band.ReadAsArray.return_value = read_data return band def _make_mock_dataset(bands_data, n_bands=None): if n_bands is None: n_bands = len(bands_data) ds = MagicMock() ds.RasterCount = n_bands ds.GetRasterBand.side_effect = lambda b: _make_mock_band(bands_data[b - 1]) return ds def _make_fake_module(name, attrs=None): mod = types.ModuleType(name) for k, v in (attrs or {}).items(): setattr(mod, k, v) return mod def _install_fake_modules(): """Install minimal fakes for scipy + osgeo so the module under test imports.""" import numpy as np # Fake scipy scipy = types.ModuleType("scipy") scipy_ndimage = types.ModuleType("scipy.ndimage") scipy_interp = types.ModuleType("scipy.interpolate") scipy_spatial = types.ModuleType("scipy.spatial") class _FakeTree: def __init__(self, coords): self.coords = coords def query(self, pts): # Return index 0 for all queries import numpy as np n = len(pts) return np.zeros(n, dtype=int), np.zeros(n, dtype=int) def _griddata(coords, values, pts, method="linear", fill_value=0.0): # Trivial: nearest valid value if len(coords) == 0 or len(pts) == 0: return np.zeros(len(pts), dtype=np.float32) # Just return first value for all queries (testing only) return np.full(len(pts), values[0], dtype=np.float32) class _FakeRBF: def __init__(self, coords, values, kernel=None): self.first_value = values[0] def __call__(self, pts): return np.full(len(pts), self.first_value, dtype=np.float32) scipy_spatial.cKDTree = _FakeTree scipy_interp.griddata = _griddata scipy_interp.RBFInterpolator = _FakeRBF scipy.ndimage = scipy_ndimage scipy.interpolate = scipy_interp scipy.spatial = scipy_spatial # Fake osgeo.gdal with the constants the module needs osgeo = types.ModuleType("osgeo") gdal_mod = types.ModuleType("osgeo.gdal") gdal_mod.GA_ReadOnly = 0 gdal_mod.GDT_Float32 = 6 gdal_mod.UseExceptions = MagicMock() def _open(path, mode): return _make_mock_dataset([]) # empty dataset; real tests build their own gdal_mod.Open = _open gdal_mod.GetDriverByName = MagicMock(return_value=MagicMock()) gdal_mod.SetConfigOption = MagicMock() osgeo.gdal = gdal_mod sys.modules["scipy"] = scipy sys.modules["scipy.ndimage"] = scipy_ndimage sys.modules["scipy.interpolate"] = scipy_interp sys.modules["scipy.spatial"] = scipy_spatial sys.modules["osgeo"] = osgeo sys.modules["osgeo.gdal"] = gdal_mod return { "scipy": scipy, "scipy.spatial": scipy_spatial, "scipy.interpolate": scipy_interp, "osgeo.gdal": gdal_mod, } class TestProcessOneBlockMocked(unittest.TestCase): """Verify _process_one_block logic with mocked GDAL/scipy.""" @classmethod def setUpClass(cls): _install_fake_modules() # Import the module under test after installing fakes from src.core.algorithms.interpolation import interpolator cls.interp = interpolator cls.gdal = sys.modules["osgeo.gdal"] def _build_dataset(self, bands): """Build a mock dataset where GetRasterBand(b).ReadAsArray(x,y,w,h) returns bands[b-1] (a numpy array of the requested shape). For simplicity we return the full band array regardless of x,y,w,h. """ ds = MagicMock() ds.RasterCount = len(bands) def get_band(b): band = MagicMock() band.ReadAsArray.return_value = bands[b - 1] return band ds.GetRasterBand.side_effect = get_band return ds def test_no_zeros_returns_inner_blocks_unchanged(self): """If no pixels are all-zero, inner blocks should be returned as-is.""" import numpy as np # 2x2 image, 3 bands, all 1s (no zeros anywhere) band = np.ones((4, 4), dtype=np.float32) # 2 inner + 2 halo bands = [band, band * 2, band * 3] ds = self._build_dataset(bands) inner_bands, zero_count = self.interp._process_one_block( dataset=ds, x0=0, y0=0, ey0=0, ex0=0, ey1=4, ex1=4, row_offset=0, col_offset=0, inner_h=2, inner_w=2, mask_segment_ext=None, method="nearest", ) self.assertEqual(zero_count, 0) self.assertEqual(len(inner_bands), 3) for ib, expected_band in zip(inner_bands, bands): self.assertEqual(ib.shape, (2, 2)) np.testing.assert_array_equal(ib, expected_band[:2, :2]) def test_with_zero_pixel_triggers_interpolation(self): """If a pixel is all-zero, _interpolate_single_band should be called.""" import numpy as np # 4x4 image, 3 bands. Top-left 2x2 is zeros, rest is 1s. band1 = np.array([ [0, 0, 1, 1], [0, 0, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], ], dtype=np.float32) band2 = band1 * 2 band3 = band1 * 3 ds = self._build_dataset([band1, band2, band3]) # Process the top-left 2x2 block (with halo, 4x4 covers all) inner_bands, zero_count = self.interp._process_one_block( dataset=ds, x0=0, y0=0, ey0=0, ex0=0, ey1=4, ex1=4, row_offset=0, col_offset=0, inner_h=2, inner_w=2, mask_segment_ext=None, method="nearest", ) self.assertGreater(zero_count, 0, "should detect zero pixels") self.assertEqual(len(inner_bands), 3) # Each inner band should be 2x2 for ib in inner_bands: self.assertEqual(ib.shape, (2, 2)) # Our fake nearest interpolation returns valid_values[0]=1 for band1 # So the inner block should be filled with non-zero values self.assertTrue(np.all(inner_bands[0] > 0)) class TestWorkerFunctionMocked(unittest.TestCase): """Verify _interpolate_block_worker uses module-global _worker_dataset.""" @classmethod def setUpClass(cls): _install_fake_modules() from src.core.algorithms.interpolation import interpolator cls.interp = interpolator def test_worker_uses_module_global_dataset(self): import numpy as np # Set the module global band = np.array([ [0, 0, 1, 1], [0, 0, 1, 1], [1, 1, 1, 1], [1, 1, 1, 1], ], dtype=np.float32) ds = self._build_dataset([band, band * 2]) self.interp._worker_dataset = ds try: task = ( 0, 0, 0, 0, 4, 4, # x0, y0, ey0, ex0, ey1, ex1 0, 0, 2, 2, # row_offset, col_offset, inner_h, inner_w None, "nearest", # mask_segment_ext, method ) x0, y0, inner_bands, zero_count, error = self.interp._interpolate_block_worker(task) self.assertIsNone(error) self.assertEqual((x0, y0), (0, 0)) self.assertGreater(zero_count, 0) self.assertIsNotNone(inner_bands) self.assertEqual(len(inner_bands), 2) finally: self.interp._worker_dataset = None def test_worker_returns_error_if_dataset_uninitialized(self): self.interp._worker_dataset = None task = (0, 0, 0, 0, 2, 2, 0, 0, 2, 2, None, "nearest") x0, y0, inner_bands, zero_count, error = self.interp._interpolate_block_worker(task) self.assertIsNotNone(error) self.assertIn("not initialized", error) self.assertIsNone(inner_bands) def _build_dataset(self, bands): ds = MagicMock() ds.RasterCount = len(bands) def get_band(b): band = MagicMock() band.ReadAsArray.return_value = bands[b - 1] return band ds.GetRasterBand.side_effect = get_band return ds class TestInitWorkerMocked(unittest.TestCase): """Verify _init_worker sets module global and config option.""" @classmethod def setUpClass(cls): _install_fake_modules() from src.core.algorithms.interpolation import interpolator cls.interp = interpolator cls.gdal_mod = sys.modules["osgeo.gdal"] def setUp(self): self.interp._worker_dataset = None def tearDown(self): self.interp._worker_dataset = None def test_init_worker_opens_and_caches(self): fake_ds = MagicMock() with patch.object(self.gdal_mod, "Open", return_value=fake_ds) as open_mock, \ patch.object(self.gdal_mod, "SetConfigOption") as cfg_mock: self.interp._init_worker("/fake/path.bsq") self.assertIs(self.interp._worker_dataset, fake_ds) open_mock.assert_called_once_with("/fake/path.bsq", 0) # Should have set GDAL_NUM_THREADS=1 cfg_mock.assert_any_call("GDAL_NUM_THREADS", "1") def test_init_worker_raises_if_open_fails(self): with patch.object(self.gdal_mod, "Open", return_value=None): with self.assertRaises(RuntimeError): self.interp._init_worker("/bad/path.bsq") if __name__ == "__main__": unittest.main(verbosity=2)