From 3fc789a7eeb9efa6255e331ab7dda03097c451d6 Mon Sep 17 00:00:00 2001 From: Lincoln Stein Date: Wed, 2 Aug 2023 18:31:10 -0400 Subject: [PATCH] fix unit tests --- tests/nodes/test_graph_execution_state.py | 9 ++++++--- tests/nodes/test_invoker.py | 9 ++++++--- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/tests/nodes/test_graph_execution_state.py b/tests/nodes/test_graph_execution_state.py index e0ee120b54..248bc6fee1 100644 --- a/tests/nodes/test_graph_execution_state.py +++ b/tests/nodes/test_graph_execution_state.py @@ -16,6 +16,7 @@ from invokeai.app.invocations.baseinvocation import ( from invokeai.app.invocations.collections import RangeInvocation from invokeai.app.invocations.math import AddInvocation, MultiplyInvocation from invokeai.app.services.invocation_services import InvocationServices +from invokeai.app.services.invocation_stats import InvocationStatsService from invokeai.app.services.graph import ( Graph, CollectInvocation, @@ -41,6 +42,9 @@ def simple_graph(): @pytest.fixture def mock_services() -> InvocationServices: # NOTE: none of these are actually called by the test invocations + graph_execution_manager = SqliteItemStorage[GraphExecutionState]( + filename=sqlite_memory, table_name="graph_executions" + ) return InvocationServices( model_manager=None, # type: ignore events=TestEventService(), @@ -51,9 +55,8 @@ def mock_services() -> InvocationServices: board_images=None, # type: ignore queue=MemoryInvocationQueue(), graph_library=SqliteItemStorage[LibraryGraph](filename=sqlite_memory, table_name="graphs"), - graph_execution_manager=SqliteItemStorage[GraphExecutionState]( - filename=sqlite_memory, table_name="graph_executions" - ), + graph_execution_manager=graph_execution_manager, + performance_statistics=InvocationStatsService(graph_execution_manager), processor=DefaultInvocationProcessor(), configuration=None, # type: ignore ) diff --git a/tests/nodes/test_invoker.py b/tests/nodes/test_invoker.py index 8eba6d468f..5985c7e8bb 100644 --- a/tests/nodes/test_invoker.py +++ b/tests/nodes/test_invoker.py @@ -11,6 +11,7 @@ from invokeai.app.services.processor import DefaultInvocationProcessor from invokeai.app.services.sqlite import SqliteItemStorage, sqlite_memory from invokeai.app.services.invoker import Invoker from invokeai.app.services.invocation_services import InvocationServices +from invokeai.app.services.invocation_stats import InvocationStatsService from invokeai.app.services.graph import ( Graph, GraphExecutionState, @@ -34,6 +35,9 @@ def simple_graph(): @pytest.fixture def mock_services() -> InvocationServices: # NOTE: none of these are actually called by the test invocations + graph_execution_manager = SqliteItemStorage[GraphExecutionState]( + filename=sqlite_memory, table_name="graph_executions" + ) return InvocationServices( model_manager=None, # type: ignore events=TestEventService(), @@ -44,10 +48,9 @@ def mock_services() -> InvocationServices: board_images=None, # type: ignore queue=MemoryInvocationQueue(), graph_library=SqliteItemStorage[LibraryGraph](filename=sqlite_memory, table_name="graphs"), - graph_execution_manager=SqliteItemStorage[GraphExecutionState]( - filename=sqlite_memory, table_name="graph_executions" - ), + graph_execution_manager=graph_execution_manager, processor=DefaultInvocationProcessor(), + performance_statistics=InvocationStatsService(graph_execution_manager), configuration=None, # type: ignore )