fix unit tests

run-llama · Feb 19, 2025 · d4ed8c5 · d4ed8c5
1 parent 4959bdb
commit d4ed8c5
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 17 deletions.
diff --git a/llama-index-core/tests/workflow/conftest.py b/llama-index-core/tests/workflow/conftest.py
@@ -44,4 +44,4 @@ def events():
 
 @pytest.fixture()
 def ctx():
-    return Context(workflow=Workflow())
+    return Context(workflow=DummyWorkflow())
diff --git a/llama-index-core/tests/workflow/test_decorator.py b/llama-index-core/tests/workflow/test_decorator.py
@@ -1,10 +1,9 @@
 import re
 
 import pytest
-
 from llama_index.core.workflow.decorators import step
 from llama_index.core.workflow.errors import WorkflowValidationError
-from llama_index.core.workflow.events import Event
+from llama_index.core.workflow.events import Event, StartEvent, StopEvent
 from llama_index.core.workflow.workflow import Workflow
 
 
@@ -22,12 +21,12 @@ def f(self, ev: Event) -> Event:
 def test_decorate_method():
     class TestWorkflow(Workflow):
         @step
-        def f1(self, ev: Event) -> Event:
+        def f1(self, ev: StartEvent) -> Event:
             return ev
 
         @step
-        def f2(self, ev: Event) -> Event:
-            return ev
+        def f2(self, ev: Event) -> StopEvent:
+            return StopEvent()
 
     wf = TestWorkflow()
     assert getattr(wf.f1, "__step_config")

diff --git a/llama-index-core/tests/workflow/test_service.py b/llama-index-core/tests/workflow/test_service.py
@@ -1,10 +1,9 @@
 import pytest
-
+from llama_index.core.workflow.context import Context
 from llama_index.core.workflow.decorators import step
 from llama_index.core.workflow.events import Event, StartEvent, StopEvent
-from llama_index.core.workflow.workflow import Workflow
-from llama_index.core.workflow.context import Context
 from llama_index.core.workflow.service import ServiceManager, ServiceNotFoundError
+from llama_index.core.workflow.workflow import Workflow
 
 
 class ServiceWorkflow(Workflow):
@@ -67,17 +66,15 @@ async def test_default_value_for_service():
     assert res == 84
 
 
-def test_service_manager_add():
+def test_service_manager_add(workflow):
     s = ServiceManager()
-    w = Workflow()
-    s.add("test_id", w)
-    assert s._services["test_id"] == w
+    s.add("test_id", workflow)
+    assert s._services["test_id"] == workflow
 
 
-def test_service_manager_get():
+def test_service_manager_get(workflow):
     s = ServiceManager()
-    w = Workflow()
-    s._services["test_id"] = w
-    assert s.get("test_id") == w
+    s._services["test_id"] = workflow
+    assert s.get("test_id") == workflow
     with pytest.raises(ServiceNotFoundError):
         s.get("not_found")