412 lines
15 KiB
Python
412 lines
15 KiB
Python
|
"""
|
|||
|
|
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)
|