Files
AnyImage/tests/test_addon.py
nan a503786c62 Refactor and enhance SnapForm functionality
- Introduced MODEL_NUMBERS for better model indexing.
- Improved enum_items function to cache results and handle download status.
- Updated ActiveJob class to manage temporary input paths and added displacement_via property.
- Changed IMAGE_DEPTH_PROCESS_RES_METHOD from "lower_bound_resize" to "upper_bound_resize".
- Enhanced RunJob operator to include new properties and improved error messages.
- Modified panel layout for better user experience and added runtime installation checks.
- Updated SnapFormSettings to reflect new model properties and defaults.
- Added temporary input path management in runtime.
- Refined worker functions to improve error handling and messaging.
- Expanded tests to cover new functionality and ensure backward compatibility.
- Fixed packaging script to align with versioning conventions.
2026-08-03 13:47:15 +08:00

824 lines
29 KiB
Python

import importlib
import inspect
import sys
import unittest
from pathlib import Path
from tempfile import TemporaryDirectory
from types import ModuleType, SimpleNamespace
from unittest.mock import patch
PROJECT_ROOT = Path(__file__).resolve().parents[1]
SOURCE_ROOT = PROJECT_ROOT / "src"
class AddonTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
cls.remove_modules()
fake_bpy = ModuleType("bpy")
fake_bpy.props = SimpleNamespace(
BoolProperty=lambda **options: options,
EnumProperty=lambda **options: options,
FloatProperty=lambda **options: options,
IntProperty=lambda **options: options,
PointerProperty=lambda **options: options,
StringProperty=lambda **options: options,
)
fake_bpy.types = SimpleNamespace(
AddonPreferences=object,
Operator=object,
Panel=object,
PropertyGroup=object,
Scene=type("Scene", (), {}),
STATUSBAR_HT_header=SimpleNamespace(
append=lambda draw: cls.status_bar_draws.append(draw),
remove=lambda draw: cls.status_bar_draws.remove(draw),
),
VIEW3D_MT_object_context_menu=SimpleNamespace(
append=lambda draw: cls.context_menu_draws.append(draw),
remove=lambda draw: cls.context_menu_draws.remove(draw),
),
)
cls.registered = []
cls.status_bar_draws = []
cls.context_menu_draws = []
fake_bpy.utils = SimpleNamespace(
register_class=cls.registered.append,
unregister_class=cls.registered.remove,
)
fake_bpy.context = SimpleNamespace(
preferences=SimpleNamespace(addons={}),
)
fake_bpy.app = SimpleNamespace(online_access=True)
fake_bpy.path = SimpleNamespace(abspath=lambda path: path)
sys.modules["bpy"] = fake_bpy
sys.path.insert(0, str(SOURCE_ROOT))
cls.addon = importlib.import_module("snapform")
cls.fake_bpy = fake_bpy
@classmethod
def tearDownClass(cls):
sys.path.remove(str(SOURCE_ROOT))
sys.modules.pop("bpy", None)
cls.remove_modules()
@staticmethod
def remove_modules():
for module_name in tuple(sys.modules):
if module_name == "snapform" or module_name.startswith("snapform."):
sys.modules.pop(module_name)
def tearDown(self):
self.registered.clear()
self.status_bar_draws.clear()
self.context_menu_draws.clear()
self.addon.operators._active_job = None
if hasattr(self.fake_bpy.types.Scene, "snapform_settings"):
del self.fake_bpy.types.Scene.snapform_settings
def test_uses_extension_manifest_instead_of_legacy_bl_info(self):
self.assertFalse(hasattr(self.addon, "bl_info"))
def test_registers_and_unregisters_extension_types(self):
self.addon.register()
self.assertEqual(self.registered, list(self.addon.CLASSES))
self.assertTrue(
hasattr(self.fake_bpy.types.Scene, "snapform_settings")
)
self.assertEqual(
self.status_bar_draws,
[self.addon.draw_status_bar],
)
self.assertEqual(
self.context_menu_draws,
[self.addon.draw_image_context_menu],
)
self.addon.unregister()
self.assertEqual(self.registered, [])
self.assertFalse(
hasattr(self.fake_bpy.types.Scene, "snapform_settings")
)
self.assertEqual(self.status_bar_draws, [])
self.assertEqual(self.context_menu_draws, [])
def test_registers_depth_plane_operator_without_cache_rebuild(self):
self.assertIn(
self.addon.ConvertImageToDepthPlane,
self.addon.CLASSES,
)
self.assertIn(self.addon.BuildDepthPlane, self.addon.CLASSES)
self.assertEqual(
self.addon.BuildDepthPlane.bl_options,
{"INTERNAL"},
)
self.assertFalse(hasattr(self.addon, "ImportDepthPlane"))
def test_completed_depth_job_uses_undoable_build_operator(self):
source = inspect.getsource(
self.addon.RunJob._build_depth_plane
)
self.assertIn("bpy.ops.snapform.build_depth_plane(", source)
self.assertIn('if "FINISHED" not in result:', source)
self.assertIn(
'bpy.ops.ed.undo_push(message="Convert to Depth Plane")',
source,
)
self.assertIn('if "FINISHED" not in undo_result:', source)
def test_image_context_menu_only_exposes_conversion(self):
draw_source = inspect.getsource(
self.addon.draw_image_context_menu
)
self.assertIn(
'"snapform.convert_image_to_depth_plane"',
draw_source,
)
self.assertIn('operator_context = "INVOKE_DEFAULT"', draw_source)
self.assertNotIn("import_depth_plane", draw_source)
self.assertNotIn("Rebuild", draw_source)
def test_conversion_dialog_exposes_requested_settings(self):
operator = self.addon.ConvertImageToDepthPlane
self.assertNotIn(
"default",
operator.__annotations__["snapform_model"],
)
self.assertTrue(operator.__annotations__["snapform_use_ray_pose"]["default"])
self.assertEqual(
operator.__annotations__["snapform_process_res_factor"]["default"],
0.5,
)
self.assertTrue(
operator.__annotations__["snapform_delete_original"]["default"]
)
self.assertEqual(
operator.__annotations__["snapform_subdivision_level"]["default"],
5,
)
self.assertEqual(
set(operator.__annotations__),
{
"snapform_depth_source",
"snapform_displacement_via",
"snapform_model",
"snapform_use_ray_pose",
"snapform_process_res_factor",
"snapform_delete_original",
"snapform_subdivision_level",
},
)
depth_sources = operator.__annotations__["snapform_depth_source"]
self.assertEqual(depth_sources["default"], "ORIGINAL")
self.assertEqual(
[item[0] for item in depth_sources["items"]],
["ESTIMATE", "ORIGINAL"],
)
self.assertEqual(depth_sources["items"][0][1], "AI Generated")
self.assertEqual(depth_sources["name"], "Depth from")
displacement_via = operator.__annotations__[
"snapform_displacement_via"
]
self.assertEqual(displacement_via["default"], "MODIFIER")
self.assertEqual(displacement_via["name"], "Displace via")
self.assertEqual(
[item[0] for item in displacement_via["items"]],
["MODIFIER", "MATERIAL"],
)
draw_source = inspect.getsource(operator.draw)
self.assertIn(
'if self.snapform_depth_source == "ESTIMATE":',
draw_source,
)
self.assertLess(
draw_source.index(
'"snapform_process_res_factor",'
),
draw_source.index(
'layout.prop(self, "snapform_subdivision_level")'
),
)
self.assertIn(
'text=f"Resolution Scale ({process_res} px)"',
draw_source,
)
self.assertIn("estimate_box = layout.box()", draw_source)
self.assertIn(
'estimate_box.prop(self, "snapform_model")',
draw_source,
)
self.assertIn(
'estimate_box.prop(self, "snapform_use_ray_pose")',
draw_source,
)
self.assertNotIn("layout.separator", draw_source)
self.assertNotIn("layout.label", draw_source)
model_items = operator.__annotations__["snapform_model"]["items"]
model_items = model_items(None, None)
self.assertEqual(
[item[0] for item in model_items],
[
"SMALL",
"BASE",
"LARGE_1_1",
"GIANT_1_1",
"NESTED_GIANT_LARGE_1_1",
],
)
self.assertIn("invoke_props_dialog", inspect.getsource(operator.invoke))
self.assertNotIn(
'"DA3"',
inspect.getsource(operator),
)
def test_depth_conversion_progress_is_provider_neutral(self):
starting_source = inspect.getsource(
self.addon.RunJob._starting_message
)
finish_source = inspect.getsource(self.addon.RunJob._finish)
self.assertIn('return "Starting depth estimation"', starting_source)
self.assertIn('"Depth plane is ready"', finish_source)
self.assertNotIn("DA3 image depth", starting_source)
self.assertNotIn("DA3 depth plane", finish_source)
def test_model_options_mark_only_models_that_are_not_downloaded(self):
items_callback = self.addon.SnapFormSettings.__annotations__["model"][
"items"
]
with patch.object(
self.addon.runtime,
"model_ready",
side_effect=lambda key: key == "BASE",
):
items = items_callback(None, None)
labels = {item[0]: item[1] for item in items}
self.assertEqual(labels["BASE"], "DA3 Base")
self.assertEqual(labels["SMALL"], "DA3 Small (Missing)")
def test_packed_image_is_exported_without_unpacking(self):
operator = self.addon.ConvertImageToDepthPlane()
operator.snapform_depth_source = "ESTIMATE"
operator.snapform_model = "BASE"
operator.snapform_use_ray_pose = True
operator.snapform_process_res_factor = 0.5
operator.snapform_displacement_via = "MATERIAL"
operator.snapform_subdivision_level = 5
operator.snapform_delete_original = True
image = SimpleNamespace(
filepath="",
packed_file=object(),
size=(1920, 1080),
)
context = SimpleNamespace(
object=SimpleNamespace(name="Reference", data=image),
scene=object(),
)
calls = []
self.fake_bpy.ops = SimpleNamespace(
snapform=SimpleNamespace(
run_job=lambda **options: calls.append(options)
or {"RUNNING_MODAL"},
),
)
with TemporaryDirectory() as directory:
packed_path = Path(directory) / "packed.png"
packed_path.touch()
with patch.object(
self.addon.image_depth,
"save_packed_image",
return_value=packed_path,
) as save_packed_image:
result = operator.execute(context)
self.assertEqual(result, {"FINISHED"})
save_packed_image.assert_called_once_with(image, context.scene)
self.assertEqual(calls[0]["snapform_input_path"], str(packed_path))
self.assertEqual(calls[0]["snapform_process_res"], 960)
self.assertEqual(
calls[0]["snapform_displacement_via"],
"MATERIAL",
)
self.assertTrue(calls[0]["snapform_delete_input"])
def test_packed_image_export_writes_a_png_copy(self):
class PackedImage:
def save_render(self, filepath, scene):
self.filepath = filepath
self.scene = scene
self.file_format = scene.render.image_settings.file_format
Path(filepath).write_bytes(b"packed image")
image = PackedImage()
image_settings = SimpleNamespace(file_format="OPEN_EXR")
scene = SimpleNamespace(
render=SimpleNamespace(image_settings=image_settings),
)
with TemporaryDirectory() as directory:
output_path = Path(directory) / "input.png"
with patch.object(
self.addon.runtime,
"temporary_input_path",
return_value=output_path,
):
result = self.addon.image_depth.save_packed_image(image, scene)
self.assertEqual(result, output_path)
self.assertEqual(output_path.read_bytes(), b"packed image")
self.assertEqual(image.filepath, str(output_path))
self.assertIs(image.scene, scene)
self.assertEqual(image.file_format, "PNG")
self.assertEqual(image_settings.file_format, "OPEN_EXR")
def test_temporary_packed_input_is_removed_when_job_closes(self):
closed = []
with TemporaryDirectory() as directory:
input_path = Path(directory) / "packed.png"
input_path.touch()
self.addon.operators._active_job = SimpleNamespace(
log_file=SimpleNamespace(close=lambda: closed.append(True)),
temporary_input_path=input_path,
)
self.addon.operators.close_active_job()
self.assertFalse(input_path.exists())
self.assertEqual(closed, [True])
self.assertIsNone(self.addon.operators.active_job())
def test_image_depth_model_and_ray_pose_are_independent(self):
settings = SimpleNamespace(
input_path="debug-input.png",
export_format="gs_ply",
model="BASE",
gs_model="GIANT_1_1",
)
operator = self.addon.RunJob()
operator.job_type = self.addon.operators.IMAGE_DEPTH_JOB
operator.snapform_source_object_name = "Reference"
operator.snapform_process_res = 336
with TemporaryDirectory() as directory:
input_path = Path(directory) / "input.png"
input_path.touch()
operator.snapform_input_path = str(input_path)
paths = {
"log": Path(directory) / "job.log",
"status": Path(directory) / "job.json",
"result": Path(directory) / "result",
}
for snapform_model, snapform_ray_pose in (
("SMALL", True),
("BASE", False),
("LARGE_1_1", True),
("GIANT_1_1", False),
("NESTED_GIANT_LARGE_1_1", True),
):
with self.subTest(model=snapform_model), patch.object(
self.addon.runtime,
"new_job_paths",
return_value=paths,
), patch.object(
self.addon.runtime,
"runtime_ready",
return_value=True,
), patch.object(
self.addon.runtime,
"model_ready",
return_value=True,
) as model_ready, patch.object(
self.addon.runtime,
"reconstruct_command",
return_value=["worker"],
) as reconstruct:
operator.snapform_model_key = snapform_model
operator.snapform_use_ray_pose = snapform_ray_pose
command, prepared_paths = operator._prepare_command(
None,
settings,
)
self.assertEqual(command, ["worker"])
self.assertEqual(prepared_paths, paths)
model_ready.assert_called_once_with(snapform_model)
self.assertEqual(reconstruct.call_args.args[1], "mini_npz")
self.assertEqual(
reconstruct.call_args.kwargs["input_path"],
input_path,
)
self.assertEqual(
reconstruct.call_args.kwargs["model_key"],
snapform_model,
)
self.assertEqual(
reconstruct.call_args.kwargs["process_res"],
336,
)
self.assertEqual(
reconstruct.call_args.kwargs["process_res_method"],
"upper_bound_resize",
)
self.assertEqual(
reconstruct.call_args.kwargs["use_ray_pose"],
snapform_ray_pose,
)
def test_preferences_use_package_identifier(self):
self.assertEqual(
self.addon.SnapFormPreferences.bl_idname,
"snapform",
)
def test_default_storage_is_shared_in_user_home(self):
runtime = importlib.import_module("snapform.runtime")
self.assertEqual(runtime.DEFAULT_STORAGE_ROOT.name, ".snapform")
self.assertEqual(
self.addon.SnapFormPreferences.__annotations__["storage_root"][
"default"
],
str(runtime.DEFAULT_STORAGE_ROOT),
)
def test_process_res_method_defaults_to_upper_bound_resize(self):
process_res = self.addon.SnapFormSettings.__annotations__[
"process_res"
]
process_res_method = (
self.addon.SnapFormSettings.__annotations__[
"process_res_method"
]
)
self.assertEqual(process_res["default"], 1024)
self.assertEqual(process_res_method["default"], "upper_bound_resize")
def test_reconstruct_command_preserves_da3_parameter_names(self):
runtime = importlib.import_module("snapform.runtime")
settings = SimpleNamespace(
input_path="input.png",
model="GIANT_1_1",
gs_model="NESTED_GIANT_LARGE_1_1",
process_res=504,
process_res_method="upper_bound_resize",
conf_thresh_percentile=40.0,
num_max_points=1_000_000,
show_cameras=True,
gs_views_interval=1,
video_frames=16,
device="CUDA",
use_ray_pose=False,
ref_view_strategy="saddle_balanced",
)
with patch.object(
runtime,
"runtime_python",
return_value=Path("python"),
), patch.object(
runtime,
"model_directory",
return_value=Path("model"),
):
command = runtime.reconstruct_command(
settings,
"gs_ply",
Path("status.json"),
Path("result"),
)
self.assertIn("--export-format", command)
self.assertIn("gs_ply", command)
for parameter in (
"--process-res",
"--process-res-method",
"--conf-thresh-percentile",
"--num-max-points",
"--show-cameras",
"--gs-views-interval",
"--use-ray-pose",
"--ref-view-strategy",
):
with self.subTest(parameter=parameter):
self.assertIn(parameter, command)
def test_reconstruct_command_accepts_image_depth_overrides(self):
runtime = importlib.import_module("snapform.runtime")
settings = SimpleNamespace(
input_path="input.png",
model="BASE",
gs_model="GIANT_1_1",
process_res=504,
process_res_method="upper_bound_resize",
conf_thresh_percentile=40.0,
num_max_points=1_000_000,
show_cameras=True,
gs_views_interval=1,
video_frames=16,
device="AUTO",
use_ray_pose=False,
ref_view_strategy="saddle_balanced",
)
with patch.object(
runtime,
"runtime_python",
return_value=Path("python"),
), patch.object(
runtime,
"model_directory",
return_value=Path("large-model"),
) as snapform_model_directory:
snapform_command = runtime.reconstruct_command(
settings,
"mini_npz",
Path("status.json"),
Path("result"),
model_key="BASE",
process_res=252,
process_res_method="lower_bound_resize",
use_ray_pose=False,
)
snapform_model_directory.assert_called_once_with("BASE")
snapform_process_res_index = snapform_command.index("--process-res")
self.assertEqual(snapform_command[snapform_process_res_index + 1], "252")
snapform_process_res_method_index = snapform_command.index("--process-res-method")
self.assertEqual(
snapform_command[snapform_process_res_method_index + 1],
"lower_bound_resize",
)
snapform_ray_pose_index = snapform_command.index("--use-ray-pose")
self.assertEqual(snapform_command[snapform_ray_pose_index + 1], "0")
def test_export_format_property_includes_all_da3_formats(self):
export_format = self.addon.SnapFormSettings.__annotations__[
"export_format"
]
identifiers = [item[0] for item in export_format["items"]]
self.assertEqual(
identifiers,
["glb", "gs_ply", "depth_vis", "mini_npz"],
)
self.assertEqual(export_format["default"], "glb")
def test_gs_model_property_only_includes_gaussian_models(self):
gs_model = self.addon.SnapFormSettings.__annotations__[
"gs_model"
]
identifiers = [item[0] for item in gs_model["items"](None, None)]
self.assertEqual(
identifiers,
["GIANT_1_1", "NESTED_GIANT_LARGE_1_1"],
)
def test_debug_panel_is_closed_by_default(self):
self.assertIn(
"DEFAULT_CLOSED",
self.addon.DebugPanel.bl_options,
)
def test_debug_panel_uses_one_export_button_and_dynamic_settings(self):
draw_source = inspect.getsource(
self.addon.DebugPanel.draw
)
self.assertEqual(draw_source.count("export_row.operator("), 1)
self.assertIn('settings.export_format == "glb"', draw_source)
self.assertIn('settings.export_format == "gs_ply"', draw_source)
self.assertIn('"gs_model" if settings.export_format', draw_source)
self.assertIn("layout.prop(settings, model_property)", draw_source)
self.assertNotIn("model_row", draw_source)
self.assertIn("export_row.enabled = active_job() is None", draw_source)
self.assertNotIn('text="Ready"', draw_source)
self.assertNotIn('text="Missing"', draw_source)
self.assertNotIn("License:", draw_source)
self.assertIn('operator.job_type = "RECONSTRUCT"', draw_source)
def test_runtime_models_panel_assigns_each_download_target(self):
draw_source = inspect.getsource(
self.addon.RuntimeModelsPanel.draw
)
model_source = inspect.getsource(
self.addon.RuntimeModelsPanel._draw_model
)
self.assertIn("for model in MODELS.values()", draw_source)
self.assertIn("operator.model_key = model.key", model_source)
self.assertIn(
"download_row.enabled = runtime_ready() and not running",
model_source,
)
self.assertNotIn('label(text="Not downloaded"', model_source)
self.assertNotIn("Isolated Runtime ready", draw_source)
self.assertNotIn("License:", model_source)
self.assertEqual(
self.addon.RuntimeModelsPanel.bl_label,
"Setup",
)
def test_preferences_install_button_is_disabled_while_busy(self):
draw_source = inspect.getsource(self.addon.SnapFormPreferences.draw)
self.assertIn(
"install_row.enabled = active_job() is None",
draw_source,
)
def test_completed_export_opens_job_directory(self):
calls = []
self.fake_bpy.ops = SimpleNamespace(
wm=SimpleNamespace(
path_open=lambda **options: calls.append(options) or {"FINISHED"},
),
)
with TemporaryDirectory() as directory:
self.addon.operators.open_job_directory(directory)
self.assertEqual(len(calls), 1)
self.assertIn("filepath", calls[0])
def test_status_bar_draws_progress_and_cancel_button(self):
progress_calls = []
operator_calls = []
class Layout:
ui_units_x = 0.0
def row(self, align=False):
return self
def label(self, **_options):
return None
def progress(self, **options):
progress_calls.append(options)
def operator(self, operator, **options):
operator_calls.append((operator, options))
self.addon.operators._active_job = SimpleNamespace()
context = SimpleNamespace(
scene=SimpleNamespace(
snapform_settings=SimpleNamespace(
status="Processing",
progress=0.25,
),
),
)
self.addon.draw_status_bar(
SimpleNamespace(layout=Layout()),
context,
)
self.assertEqual(
progress_calls,
[{"factor": 0.25, "type": "BAR", "text": "Processing"}],
)
self.assertEqual(
operator_calls,
[("snapform.cancel_job", {"text": "", "icon": "X"})],
)
def test_image_depth_status_uses_same_progress_layout(self):
progress_calls = []
operator_calls = []
class Layout:
def row(self, align=False):
return self
def label(self, **_options):
return None
def progress(self, **options):
progress_calls.append(options)
def operator(self, operator, **options):
operator_calls.append((operator, options))
self.addon.operators._active_job = SimpleNamespace(
kind=self.addon.operators.IMAGE_DEPTH_JOB,
)
context = SimpleNamespace(
scene=SimpleNamespace(
snapform_settings=SimpleNamespace(
status="Estimating depth",
progress=0.25,
),
),
)
self.addon.draw_status_bar(
SimpleNamespace(layout=Layout()),
context,
)
self.assertEqual(
progress_calls,
[
{
"factor": 0.25,
"type": "BAR",
"text": "Estimating depth",
}
],
)
self.assertEqual(
operator_calls,
[("snapform.cancel_job", {"text": "", "icon": "X"})],
)
def test_progress_is_appended_before_cancel_button(self):
register_source = inspect.getsource(self.addon.register)
draw_source = inspect.getsource(self.addon.draw_status_bar)
self.assertIn("STATUSBAR_HT_header.append", register_source)
self.assertNotIn("STATUSBAR_HT_header.prepend", register_source)
self.assertLess(
draw_source.index("progress_row.progress("),
draw_source.index('"snapform.cancel_job"'),
)
self.assertIn("progress_row.ui_units_x = 12.0", draw_source)
self.assertNotIn('label(text="SnapForm"', draw_source)
def test_ui_redraw_targets_status_bar_and_view_3d_areas(self):
redrawn = []
status_regions = (
SimpleNamespace(
tag_redraw=lambda: redrawn.append("status-region"),
),
)
areas = (
SimpleNamespace(
type="STATUSBAR",
tag_redraw=lambda: redrawn.append("status"),
regions=status_regions,
),
SimpleNamespace(
type="VIEW_3D",
tag_redraw=lambda: redrawn.append("view"),
regions=(),
),
)
context = SimpleNamespace(
window_manager=SimpleNamespace(
windows=(
SimpleNamespace(screen=SimpleNamespace(areas=areas)),
),
),
)
self.addon.status_bar.redraw_ui(context)
self.addon.status_bar.redraw_ui(context)
self.assertEqual(
redrawn,
[
"status",
"status-region",
"view",
"status",
"status-region",
"view",
],
)
def test_force_flag_does_not_trigger_window_draw_and_swap(self):
redrawn = []
context = SimpleNamespace(
window_manager=SimpleNamespace(
windows=(
SimpleNamespace(
screen=SimpleNamespace(
areas=(
SimpleNamespace(
type="STATUSBAR",
tag_redraw=lambda: redrawn.append(
"status"
),
regions=(),
),
),
),
),
),
),
)
self.addon.status_bar.redraw_ui(context, force=True)
self.assertEqual(redrawn, ["status"])
source = inspect.getsource(self.addon.status_bar)
self.assertNotIn("redraw_timer", source)
self.assertNotIn("DRAW_WIN_SWAP", source)