test(workflow): add tests for workflow feature

3 files changed, 55 insertions, 12 deletions

diff --git a/tests/test_core.py b/tests/test_core.py
index 5f3f12f1..bcab2d4a 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -1,5 +1,5 @@
 from infini.core import Core
-from infini.generator import TextGenerator
+from infini.generator import Generator
 from infini.handler import Handler
 from infini.injector import Injector
 from infini.input import Input
@@ -28,12 +28,10 @@ def test_core():
 
     def add(input: Input) -> Output:
         result = str(sum(list(map(int, input.get_plain_text().lstrip(".add").split()))))
-        return Output(
-            "text", "test.add", status=0, block=False, variables={"result": result}
-        )
+        return Output("text", "test.add", block=False, variables={"result": result})
 
     def cmd(_: Input) -> Output:
-        return Output("text", "test.cmd", status=0, block=False)
+        return Output("text", "test.cmd", block=False)
 
     handler = Handler()
     handler.handlers = [
@@ -49,12 +47,13 @@ def test_core():
         },
     ]
 
-    generator = TextGenerator()
+    generator = Generator()
     generator.events = {
         "test.cmd": "cmd",
         "test.add": "{{ result }}",
         "block.jianlvchun": "检测到违禁词",
     }
+    generator.global_variables = {}
 
     core = Core()
     core.handler = handler
diff --git a/tests/test_handlers.py b/tests/test_handlers.py
index ee96cdab..79810ce0 100644
--- a/tests/test_handlers.py
+++ b/tests/test_handlers.py
@@ -11,12 +11,11 @@ def test_handler():
         return Output(
             "text",
             str(sum(list(map(int, input.get_plain_text().lstrip(".add").split())))),
-            status=0,
             block=False,
         )
 
     def cmd(_: Input) -> Output:
-        return Output("text", "cmd", status=0, block=False)
+        return Output("text", "cmd", block=False)
 
     handler = Handler()
     handler.handlers = [
@@ -45,12 +44,11 @@ def test_handler_block():
         return Output(
             "text",
             str(sum(list(map(int, input.get_plain_text().lstrip(".add").split())))),
-            status=0,
             block=False,
         )
 
     def cmd(_: Input) -> Output:
-        return Output("text", "cmd", status=0, block=True)
+        return Output("text", "cmd", block=True)
 
     handler = Handler()
     handler.handlers = [
@@ -79,10 +77,9 @@ def test_handler_interator():
         yield Output(
             "text",
             str(sum(list(map(int, input.get_plain_text().lstrip(".add").split())))),
-            status=0,
             block=False,
         )
-        yield Output("text", "ok", status=0, block=False)
+        yield Output("text", "ok", block=False)
 
     handler = Handler()
     handler.handlers = [
diff --git a/tests/test_workflow.py b/tests/test_workflow.py
new file mode 100644
index 00000000..18dab9d5
--- /dev/null
+++ b/tests/test_workflow.py
@@ -0,0 +1,47 @@
+from infini.core import Core
+from infini.generator import Generator
+from infini.handler import Handler
+from infini.injector import Injector
+from infini.input import Input
+from infini.interceptor import Interceptor
+from infini.output import Output
+from infini.router import Startswith
+
+
+def test_workflow():
+    def func_workflow(input: Input):
+        yield input.output("workflow", "test.workflow", block=True)
+
+    input = Input("testmsg")
+
+    handler = Handler()
+    handler.handlers = [
+        {
+            "priority": 0,
+            "router": Startswith(""),
+            "handler": func_workflow,
+        }
+    ]
+
+    interceptor = Interceptor()
+    interceptor.interceptors = []
+
+    generator = Generator()
+    generator.events = {
+        "test.cmd": "cmd",
+        "test.add": "{{ result }}",
+        "block.jianlvchun": "检测到违禁词",
+    }
+    generator.global_variables = {}
+
+    core = Core()
+    core.handler = handler
+    core.interceptor = interceptor
+    core.pre_interceptor = interceptor
+    core.generator = generator
+    core.injector = Injector()
+
+    for output in core.input(input):
+        assert isinstance(output, Output)
+        assert output.type == "workflow"
+        output.status = 0