"""Batch Export View as Images: does the exported PNG actually match what the live view would show, and does the batch dispatch (menu action -> worker -> per-file render) behave like Batch Compute's proven pattern? sras_render.export_view_image is tested directly (no Qt) for the plumbing that decides *what* gets rendered -- pad_factor derived per file, masked- pixel NaN fill, per-file auto-scale, velocity scaling -- via a draw_view_image spy rather than pixel-diffing PNGs, the same "spy on the seam, don't inspect the rendered artifact" approach test_highlight_masked_ pixels (tests/test_gui.py) uses for the live canvas. The GUI-dispatch half drives SrasViewerWindow._on_batch_export_images() end-to-end with patched file dialogs, the same shape as test_stored_cache.py's test_viewer_batch_row_average_dispatch. """ from pathlib import Path from unittest.mock import patch import numpy as np import pytest from PyQt6.QtCore import QEventLoop, QTimer from PyQt6.QtWidgets import QApplication import sras_render from sras_compute import compute_rf_image, dc_image_mv from sras_format import CH1_IDX, CH3_IDX, CH4_IDX, CH_NAMES, SrasFile from sras_render import export_view_image from sras_viewer import SrasViewerWindow, VELOCITY_MODE_IDX from sras_workers import BatchExportImagesWorker import tools.make_test_sras as gen _THRESHOLD_MV = 50.0 def pump(ms: int = 200): loop = QEventLoop() QTimer.singleShot(ms, loop.quit) loop.exec() def wait_until(pred, timeout_ms: int = 20000, step: int = 100) -> bool: waited = 0 while waited < timeout_ms: if pred(): return True pump(step) waited += step return pred() # --------------------------------------------------------------------------- # sras_render.export_view_image -- pure function, no Qt # --------------------------------------------------------------------------- _DEFAULT_KW = dict( is_fft_mode=False, is_velocity=False, dc_threshold_mv=_THRESHOLD_MV, apply_bg_sub=True, pad_factor=1, min_freq_mhz=0.0, grating_um=1.0, cmap="viridis", auto_scale=True, vmin=0.0, vmax=1.0, highlight_masked=False, mode_str="DC", colorbar_label="mV", ) def _kw(**overrides): kw = dict(_DEFAULT_KW) kw.update(overrides) return kw def _spy_draw(monkeypatch): """Patches sras_render.draw_view_image to record the image array and vmin/vmax/bad_color it was called with, then delegates to the real implementation so the PNG is still written -- lets a test check *what* export_view_image computed without depending on rendered PNG pixels.""" orig = sras_render.draw_view_image captured = {} def spy(ax, fig, img, extent, cmap, vmin, vmax, xlabel, ylabel, title, colorbar_label="", cb_ticks=None, norm=None, bad_color=None): captured["img"] = np.array(img, copy=True) captured["vmin"] = vmin captured["vmax"] = vmax captured["bad_color"] = bad_color return orig(ax, fig, img, extent, cmap, vmin, vmax, xlabel, ylabel, title, colorbar_label, cb_ticks, norm, bad_color) monkeypatch.setattr(sras_render, "draw_view_image", spy) return captured def test_export_dc_channel_writes_png(tmp_path): path = tmp_path / "dc.sras" gen.write(path, n_angles=2, seed=1, samples_per_frame=64) out_dir = tmp_path / "out" out_dir.mkdir() err, out_name = export_view_image( str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH4_IDX, **_kw()) assert err == "" assert out_name == f"dc_angle0_{CH_NAMES[CH4_IDX]}.png" out_path = out_dir / out_name assert out_path.exists() and out_path.stat().st_size > 0 assert out_path.read_bytes()[:8] == b"\x89PNG\r\n\x1a\n", "not a valid PNG" def test_export_out_of_range_angle(tmp_path): path = tmp_path / "short.sras" gen.write(path, n_angles=2, seed=2, samples_per_frame=64) out_dir = tmp_path / "out" out_dir.mkdir() err, out_name = export_view_image( str(path), out_dir=str(out_dir), angle_idx=5, ch_idx=CH4_IDX, **_kw()) assert err != "" and "2" in err, "error should mention the file's actual angle count" assert out_name == "" assert list(out_dir.iterdir()) == [], "no file written for a failed export" def test_export_fft_mode_matches_compute_rf_image(tmp_path, monkeypatch): path = tmp_path / "fft.sras" gen.write(path, n_angles=1, seed=3, samples_per_frame=128) out_dir = tmp_path / "out" out_dir.mkdir() captured = _spy_draw(monkeypatch) err, _ = export_view_image( str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX, **_kw(is_fft_mode=True)) assert err == "" sras = SrasFile(str(path)) expected = compute_rf_image(sras, 0, dc_threshold_mv=_THRESHOLD_MV, apply_bg_sub=True, n_fft=None, min_freq_mhz=0.0) assert np.array_equal(captured["img"], expected) def test_export_velocity_scales_frequency(tmp_path, monkeypatch): path = tmp_path / "vel.sras" gen.write(path, n_angles=1, seed=4, samples_per_frame=128) out_dir = tmp_path / "out" out_dir.mkdir() captured = _spy_draw(monkeypatch) grating_um = 3.5 err, _ = export_view_image( str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=VELOCITY_MODE_IDX, **_kw(is_fft_mode=True, is_velocity=True, grating_um=grating_um)) assert err == "" sras = SrasFile(str(path)) freq = compute_rf_image(sras, 0, dc_threshold_mv=_THRESHOLD_MV, apply_bg_sub=True, n_fft=None, min_freq_mhz=0.0) assert np.array_equal(captured["img"], freq * grating_um) def test_pad_factor_uses_each_files_own_samples_per_frame(tmp_path, monkeypatch): """n_fft must be derived per file from that file's own samples_per_frame, never a value carried over from whichever file the caller had open -- otherwise every file but one in a batch gets silently mis-padded.""" orig_draw = sras_render.draw_view_image out_dir = tmp_path / "out" out_dir.mkdir() for i, spf in enumerate((64, 256)): path = tmp_path / f"pad_{spf}.sras" gen.write(path, n_angles=1, seed=5 + i, samples_per_frame=spf) captured = {} def spy(ax, fig, img, *a, __c=captured, **kw): __c["img"] = np.array(img, copy=True) return orig_draw(ax, fig, img, *a, **kw) monkeypatch.setattr(sras_render, "draw_view_image", spy) err, _ = export_view_image( str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX, **_kw(is_fft_mode=True, pad_factor=4)) assert err == "" sras = SrasFile(str(path)) expected = compute_rf_image(sras, 0, dc_threshold_mv=_THRESHOLD_MV, apply_bg_sub=True, n_fft=spf * 4, min_freq_mhz=0.0) assert np.array_equal(captured["img"], expected), \ f"samples_per_frame={spf}: n_fft must use this file's own value" def test_highlight_masked_sets_nan_and_bad_color(tmp_path, monkeypatch): path = tmp_path / "mask.sras" gen.write(path, n_angles=1, seed=7, samples_per_frame=128) out_dir = tmp_path / "out" out_dir.mkdir() sras = SrasFile(str(path)) dc4 = dc_image_mv(sras, 0, CH4_IDX) threshold = float(np.median(dc4)) expect_masked = dc4 < threshold assert expect_masked.any() and not expect_masked.all(), \ "fixture threshold should mask some but not all pixels" captured = _spy_draw(monkeypatch) err, _ = export_view_image( str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX, **_kw(is_fft_mode=True, dc_threshold_mv=threshold, highlight_masked=True, mask_color="magenta")) assert err == "" assert captured["bad_color"] == "magenta" # The export masks by value too (0 == the "no valid peak" sentinel, same # rule as the viewer's _redraw_image); on this fixture every # above-threshold pixel has a nonzero peak, so the value mask coincides # with the DC mask and the NaN set is exactly expect_masked. assert np.array_equal(np.isnan(captured["img"]), expect_masked) valid_vals = captured["img"][~expect_masked] assert not np.isnan(valid_vals).any() and (valid_vals != 0).all(), \ "fixture precondition: every valid pixel has a nonzero peak" captured2 = _spy_draw(monkeypatch) err, _ = export_view_image( str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH1_IDX, **_kw(is_fft_mode=True, dc_threshold_mv=threshold, highlight_masked=False)) assert err == "" assert captured2["bad_color"] is None assert not np.isnan(captured2["img"]).any() def test_auto_scale_uses_per_file_min_max(tmp_path, monkeypatch): path = tmp_path / "scale.sras" gen.write(path, n_angles=1, seed=8, samples_per_frame=64) out_dir = tmp_path / "out" out_dir.mkdir() captured = _spy_draw(monkeypatch) err, _ = export_view_image( str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH4_IDX, **_kw(auto_scale=True)) assert err == "" img = captured["img"] assert captured["vmin"] == pytest.approx(float(np.nanmin(img))) assert captured["vmax"] == pytest.approx(float(np.nanmax(img))) captured2 = _spy_draw(monkeypatch) err, _ = export_view_image( str(path), out_dir=str(out_dir), angle_idx=0, ch_idx=CH4_IDX, **_kw(auto_scale=False, vmin=-5.0, vmax=5.0)) assert err == "" assert captured2["vmin"] == -5.0 assert captured2["vmax"] == 5.0 # --------------------------------------------------------------------------- # GUI dispatch: SrasViewerWindow._on_batch_export_images end-to-end # --------------------------------------------------------------------------- def _make_window(path) -> SrasViewerWindow: app = QApplication.instance() or QApplication([]) # noqa: F841 win = SrasViewerWindow() win.show() win._load_file(str(path)) assert wait_until(lambda: win._sras is not None), "file loaded" assert wait_until(lambda: all((a, CH4_IDX) in win._dc_cache for a in range(win._sras.n_angles))), \ "DC precompute finished" return win def test_batch_export_images_writes_one_png_per_file(tmp_path): path = tmp_path / "src.sras" gen.write(path, n_angles=2, seed=10, samples_per_frame=64) paths = [str(path)] for i in range(2): p2 = tmp_path / f"other{i}.sras" gen.write(p2, n_angles=2, seed=20 + i, samples_per_frame=64) paths.append(str(p2)) out_dir = tmp_path / "images" out_dir.mkdir() win = _make_window(path) try: with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames", return_value=(paths, "")), \ patch("sras_viewer.main_window.QFileDialog.getExistingDirectory", return_value=str(out_dir)): win._on_batch_export_images() assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran" angle = win.spin_angle.value() ch_name = CH_NAMES[win.combo_channel.currentIndex()] expected_names = {f"{Path(p).stem}_angle{angle}_{ch_name}.png" for p in paths} actual_names = {p.name for p in out_dir.iterdir()} assert actual_names == expected_names assert "Batch export: 3/3 image(s)" in win.statusBar().currentMessage() finally: win.close() pump(200) def test_batch_export_out_of_range_angle_reports_error_continues(tmp_path): good_path = tmp_path / "good.sras" short_path = tmp_path / "short.sras" gen.write(good_path, n_angles=3, seed=30, samples_per_frame=64) gen.write(short_path, n_angles=1, seed=31, samples_per_frame=64) out_dir = tmp_path / "images" out_dir.mkdir() win = _make_window(good_path) try: win.spin_angle.setValue(2) # valid for good_path, out of range for short_path with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames", return_value=([str(good_path), str(short_path)], "")), \ patch("sras_viewer.main_window.QFileDialog.getExistingDirectory", return_value=str(out_dir)): win._on_batch_export_images() assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran" msg = win.statusBar().currentMessage() assert "Batch export: 1/2 image(s)" in msg, msg assert "1 failed" in msg, msg assert len(list(out_dir.iterdir())) == 1, \ "the batch must not abort — the good file still exports" finally: win.close() pump(200) def test_batch_export_busy_guard_skips_dialogs(tmp_path, monkeypatch): path = tmp_path / "busy.sras" gen.write(path, n_angles=1, seed=40, samples_per_frame=64) win = _make_window(path) try: monkeypatch.setattr(win, "_job_running", lambda key: True) with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames") as mock_dlg: win._on_batch_export_images() assert mock_dlg.call_count == 0, \ "the Jobs.BATCH busy guard must return before opening any dialog" finally: win.close() pump(200) def test_batch_export_ignores_aligned_view_toggle(tmp_path): """Aligned View is geometry specific to whichever single file the Alignment Wizard last ran against and cannot be meaningfully applied across a batch of different files -- _on_batch_export_images must not read chk_aligned_view / self._alignment_result at all, regardless of what's checked in the live view.""" path = tmp_path / "aligned.sras" gen.write(path, n_angles=1, seed=41, samples_per_frame=64) out_dir = tmp_path / "images" out_dir.mkdir() win = _make_window(path) try: captured_kwargs = [] orig_init = BatchExportImagesWorker.__init__ def spy_init(self, paths, **kw): captured_kwargs.append(kw) return orig_init(self, paths, **kw) win.chk_aligned_view.setChecked(True) with patch.object(BatchExportImagesWorker, "__init__", spy_init), \ patch("sras_viewer.main_window.QFileDialog.getOpenFileNames", return_value=([str(path)], "")), \ patch("sras_viewer.main_window.QFileDialog.getExistingDirectory", return_value=str(out_dir)): win._on_batch_export_images() assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran" assert len(captured_kwargs) == 1 assert not any("align" in k.lower() for k in captured_kwargs[0]), \ captured_kwargs[0].keys() finally: win.close() pump(200) def test_batch_export_filename_collision_note(tmp_path): dir_a, dir_b = tmp_path / "dir_a", tmp_path / "dir_b" dir_a.mkdir() dir_b.mkdir() path_a, path_b = dir_a / "dup.sras", dir_b / "dup.sras" gen.write(path_a, n_angles=1, seed=50, samples_per_frame=64) gen.write(path_b, n_angles=1, seed=51, samples_per_frame=64) out_dir = tmp_path / "images" out_dir.mkdir() win = _make_window(path_a) try: with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames", return_value=([str(path_a), str(path_b)], "")), \ patch("sras_viewer.main_window.QFileDialog.getExistingDirectory", return_value=str(out_dir)): win._on_batch_export_images() assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran" msg = win.statusBar().currentMessage() assert "Batch export: 2/2 image(s)" in msg, msg assert "collision" in msg, msg assert len(list(out_dir.iterdir())) == 1, \ "same-stem inputs silently overwrite to one output file" finally: win.close() pump(200) def test_batch_export_does_not_modify_source_files(tmp_path): path = tmp_path / "untouched.sras" gen.write(path, n_angles=1, seed=60, samples_per_frame=64) before = path.read_bytes() out_dir = tmp_path / "images" out_dir.mkdir() win = _make_window(path) try: with patch("sras_viewer.main_window.QFileDialog.getOpenFileNames", return_value=([str(path)], "")), \ patch("sras_viewer.main_window.QFileDialog.getExistingDirectory", return_value=str(out_dir)): win._on_batch_export_images() assert wait_until(lambda: not win._job_running("batch"), 60000), "batch ran" finally: win.close() pump(200) assert path.read_bytes() == before, "export must never write to the source file"