wehub-resource-sync
cddb07a176
docs / deploy (push) Has been cancelled
docs / changes (push) Has been cancelled
docs / check-and-build (push) Has been cancelled
build container image / cpu (push) Has been cancelled
build container image / cuda (push) Has been cancelled
build container image / rocm (push) Has been cancelled
frontend checks / frontend-checks (push) Has been cancelled
frontend tests / frontend-tests (push) Has been cancelled
lfs checks / lfs-check (push) Has been cancelled
python checks / python-checks (push) Has been cancelled
python tests / py3.12: macos-default (push) Has been cancelled
python tests / py3.11: windows-cpu (push) Has been cancelled
python tests / py3.12: windows-cpu (push) Has been cancelled
python tests / py3.11: linux-cpu (push) Has been cancelled
typegen checks / typegen-checks (push) Has been cancelled
uv lock checks / uv-lock-checks (push) Has been cancelled
openapi checks / openapi-checks (push) Has been cancelled
python tests / py3.11: macos-default (push) Has been cancelled
python tests / py3.12: linux-cpu (push) Has been cancelled
47 lines
1.6 KiB
Python
47 lines
1.6 KiB
Python
from unittest import mock
|
|
|
|
from invokeai.backend.stable_diffusion.denoise_context import DenoiseContext
|
|
from invokeai.backend.stable_diffusion.extension_callback_type import ExtensionCallbackType
|
|
from invokeai.backend.stable_diffusion.extensions.base import ExtensionBase, callback
|
|
|
|
|
|
class MockExtension(ExtensionBase):
|
|
"""A mock ExtensionBase subclass for testing purposes."""
|
|
|
|
def __init__(self, x: int):
|
|
super().__init__()
|
|
self._x = x
|
|
|
|
@callback(ExtensionCallbackType.PRE_DENOISE_LOOP)
|
|
def set_step_index(self, ctx: DenoiseContext):
|
|
ctx.step_index = self._x
|
|
|
|
|
|
def test_extension_base_callback_registration():
|
|
"""Test that a callback can be successfully registered with an extension."""
|
|
val = 5
|
|
mock_extension = MockExtension(val)
|
|
|
|
mock_ctx = mock.MagicMock()
|
|
|
|
callbacks = mock_extension.get_callbacks()
|
|
pre_denoise_loop_cbs = callbacks.get(ExtensionCallbackType.PRE_DENOISE_LOOP, [])
|
|
assert len(pre_denoise_loop_cbs) == 1
|
|
|
|
# Call the mock callback.
|
|
pre_denoise_loop_cbs[0].function(mock_ctx)
|
|
|
|
# Confirm that the callback ran.
|
|
assert mock_ctx.step_index == val
|
|
|
|
|
|
def test_extension_base_empty_callback_type():
|
|
"""Test that an empty list is returned when no callbacks are registered for a given callback type."""
|
|
mock_extension = MockExtension(5)
|
|
|
|
# There should be no callbacks registered for POST_DENOISE_LOOP.
|
|
callbacks = mock_extension.get_callbacks()
|
|
|
|
post_denoise_loop_cbs = callbacks.get(ExtensionCallbackType.POST_DENOISE_LOOP, [])
|
|
assert len(post_denoise_loop_cbs) == 0
|