Add unit tests for service functionality and model caching

- Implement tests for ben2 and da3 model caching to ensure models are reused for matching runtime keys.
- Verify service state management, including client registration, job cancellation, and client expiration.
- Test discovery file creation and removal based on service ownership.
- Ensure frame extraction functionality works without models and produces expected results.
This commit is contained in:
2026-08-13 00:35:40 +08:00
parent 7d0983529e
commit 92d2cea1d1
21 changed files with 935 additions and 887 deletions

View File

@@ -49,7 +49,7 @@ VIDEO_NAME_SUFFIXES = (
def image_depth_model_items(owner, context):
from .models import model_enum_items
from .runtime import model_ready
from .client import model_ready
return [
item
@@ -223,7 +223,7 @@ def draw_process_resolution(layout, owner, image):
def save_packed_image(image, scene):
from .runtime import temporary_input_path
from .client import temporary_input_path
output_path = temporary_input_path()
image_settings = scene.render.image_settings
@@ -281,7 +281,7 @@ def prepare_image_input(image, scene):
return source_path, False
sequence_frames = collect_sequence_frames(image)
if sequence_frames:
from .runtime import temporary_input_directory
from .client import temporary_input_directory
temporary_directory = Path(temporary_input_directory())
try:
@@ -881,12 +881,14 @@ def material_node_group(depth_plane=False):
return existing
group = bpy.data.node_groups.new(group_name, "ShaderNodeTree")
interface = group.interface
for name, socket_type in (
input_sockets = [
("Color", "NodeSocketColor"),
("Alpha", "NodeSocketFloat"),
("Depth", "NodeSocketFloat"),
("Bump Scale", "NodeSocketFloat"),
):
]
if depth_plane:
input_sockets.insert(2, ("Depth", "NodeSocketFloat"))
for name, socket_type in input_sockets:
socket = interface.new_socket(
name=name,
in_out="INPUT",
@@ -934,7 +936,6 @@ def material_node_group(depth_plane=False):
displacement = None
processed_alpha = group_input.outputs["Alpha"]
processed_depth = group_input.outputs["Depth"]
bump_depth = group_input.outputs["Color"]
if depth_plane:
displacement = nodes.new("ShaderNodeDisplacement")
@@ -1174,8 +1175,6 @@ def create_material(
output.inputs["Displacement"],
)
set_displacement_only(material)
else:
links.new(color_texture.outputs["Color"], material_inputs.inputs["Depth"])
return material
@@ -1700,13 +1699,13 @@ class RemoveImageBackground(bpy.types.Operator):
@classmethod
def poll(cls, context):
from .operators import active_job
from .runtime import host_busy, model_ready
from .client import service_busy, model_ready
from .models import BEN2_MODEL
return (
is_image_empty(context.object)
and active_job() is None
and not host_busy()
and not service_busy()
and model_ready(BEN2_MODEL.key)
)
@@ -1790,13 +1789,13 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
@classmethod
def poll(cls, context):
from .operators import active_job
from .runtime import host_busy, model_ready
from .client import service_busy, model_ready
from .models import MODELS
return (
is_image_empty(context.object)
and active_job() is None
and not host_busy()
and not service_busy()
and any(model_ready(model_key) for model_key in MODELS)
)
@@ -1823,7 +1822,7 @@ class ConvertImageToDepthPlane(bpy.types.Operator):
available_models = image_depth_model_items(self, context)
if not available_models:
estimate_box.label(text="No DA3 models installed", icon="ERROR")
estimate_box.label(text="Download one in AI Backend > Models")
estimate_box.label(text="Download one in Models")
else:
estimate_box.prop(self, "snapform_model")
estimate_box.prop(self, "snapform_use_ray_pose")
@@ -1917,17 +1916,17 @@ class SnapFormImageMenu(bpy.types.Menu):
bl_label = "SnapForm"
def draw(self, _context):
from .runtime import runtime_ready
from .client import env_installed
layout = self.layout
layout.operator_context = "INVOKE_DEFAULT"
if not runtime_ready():
if not env_installed():
operator = layout.operator(
"snapform.run_job",
text="Install Backend",
text="Install Environment",
icon="IMPORT",
)
operator.job_type = "INSTALL_RUNTIME"
operator.job_type = "INSTALL_ENV"
return
layout.operator(
"snapform.convert_image_to_depth_plane",