58 lines
1.8 KiB
Python
58 lines
1.8 KiB
Python
# SimpleSyrup - workflow-focused ComfyUI extensions for image generation
|
|
# Copyright (C) 2026 Artificial Sweetener and contributors
|
|
# SPDX-License-Identifier: AGPL-3.0-or-later
|
|
|
|
"""Tests for the Load Ultralytics Model node contract."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any, cast
|
|
|
|
import pytest
|
|
|
|
from simple_syrup.nodes.load_ultralytics_model import LoadUltralyticsModel
|
|
from simple_syrup.runtime.ultralytics_loader import LoadedUltralyticsDetector
|
|
|
|
|
|
def test_load_ultralytics_model_node_contract(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Load Ultralytics Model exposes the planned return types and category."""
|
|
|
|
monkeypatch.setattr(LoadUltralyticsModel, "service_class", _FakeLoaderService)
|
|
|
|
inputs = LoadUltralyticsModel.INPUT_TYPES()
|
|
|
|
assert LoadUltralyticsModel.RETURN_TYPES == (
|
|
"DETECTOR_MODEL",
|
|
"BBOX_DETECTOR",
|
|
"SEGM_DETECTOR",
|
|
)
|
|
assert LoadUltralyticsModel.CATEGORY == "SimpleSyrup/Detection"
|
|
assert inputs["required"]["model_name"][0] == ["model.pt"]
|
|
|
|
|
|
def test_load_ultralytics_model_node_returns_loaded_outputs(
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
"""Load Ultralytics Model forwards native and compatibility outputs."""
|
|
|
|
monkeypatch.setattr(LoadUltralyticsModel, "service_class", _FakeLoaderService)
|
|
|
|
assert LoadUltralyticsModel().load("model.pt") == ("native", "bbox", "segm")
|
|
|
|
|
|
class _FakeLoaderService:
|
|
"""Fake loader service for node tests."""
|
|
|
|
def model_choices(self) -> list[str]:
|
|
"""Return deterministic dropdown choices."""
|
|
|
|
return ["model.pt"]
|
|
|
|
def load(self, model_name: str) -> LoadedUltralyticsDetector:
|
|
"""Return deterministic loaded outputs."""
|
|
|
|
assert model_name == "model.pt"
|
|
return LoadedUltralyticsDetector(cast(Any, "native"), "bbox", "segm")
|