From 766dc74a04bfd8f789b42d892d71ea33cc4c5a42 Mon Sep 17 00:00:00 2001
From: 苏向夜 <fu050409@163.com>
Date: Sat, 27 Jan 2024 19:08:18 +0800
Subject: :white_check_mark: test(refact): sync tests with infini changes

---
 tests/test_core.py        |  4 ++--
 tests/test_generator.py   |  8 ++++----
 tests/test_interceptor.py | 22 +++++++++++++++-------
 tests/test_register.py    |  4 ++--
 4 files changed, 23 insertions(+), 15 deletions(-)

(limited to 'tests')

diff --git a/tests/test_core.py b/tests/test_core.py
index db402dea..01d7ecb4 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -1,5 +1,5 @@
 from infini.core import Core
-from infini.generator import Generator
+from infini.generator import TextGenerator
 from infini.handler import Handler
 from infini.input import Input
 from infini.interceptor import Interceptor
@@ -48,7 +48,7 @@ def test_core():
         },
     ]
 
-    generator = Generator()
+    generator = TextGenerator()
     generator.events = {
         "test.cmd": "cmd",
         "test.add": "{{ result }}",
diff --git a/tests/test_generator.py b/tests/test_generator.py
index bcab31ed..0eed5756 100644
--- a/tests/test_generator.py
+++ b/tests/test_generator.py
@@ -1,9 +1,9 @@
-from infini.generator import Generator
+from infini.generator import TextGenerator
 from infini.output import Output
 
 
 def test_generator():
-    generator = Generator()
+    generator = TextGenerator()
     generator.events = {
         "test.event1": "Event1 文本",
     }
@@ -12,7 +12,7 @@ def test_generator():
 
 
 def test_generator_with_var():
-    generator = Generator()
+    generator = TextGenerator()
     generator.events = {
         "test.event1": "Event1 文本: {{ var }}",
     }
@@ -27,7 +27,7 @@ def test_generator_with_function():
     def add(a, b):
         return a + b
 
-    generator = Generator()
+    generator = TextGenerator()
     generator.events = {
         "test.event": "{{ func(1, 2) }}",
     }
diff --git a/tests/test_interceptor.py b/tests/test_interceptor.py
index fb12beac..afb135e9 100644
--- a/tests/test_interceptor.py
+++ b/tests/test_interceptor.py
@@ -9,7 +9,7 @@ def test_interceptor():
     valid_input = Input("这个叫苏向夜.")
 
     def intercept(_: Input) -> Input | Output:
-        return Output("text", "block.jianlvchun", block=True)  # TODO 拦截器阻塞标识
+        return Output("text", "block.jianlvchun", block=True)
 
     interceptor = Interceptor()
     interceptor.interceptors = [
@@ -19,10 +19,18 @@ def test_interceptor():
             "handler": intercept,
         }
     ]
-    output = interceptor.input(input)
-    assert isinstance(output, Output)
-    assert output.name == "block.jianlvchun"
+    for output in interceptor.input(input):
+        assert isinstance(output, Output)
+        assert output.name == "block.jianlvchun"
 
-    valid_output = interceptor.input(valid_input)
-    assert isinstance(valid_output, Input)
-    assert valid_output.get_plain_text() == "这个叫苏向夜."
+    for valid_output in interceptor.input(valid_input):
+        assert isinstance(valid_output, Input)
+        assert valid_output.get_plain_text() == "这个叫苏向夜."
+
+    for output in interceptor.output("简律纯"):
+        assert isinstance(output, Output)
+        assert output.name == "block.jianlvchun"
+
+    for output in interceptor.output("这个叫苏向夜."):
+        assert isinstance(output, str)
+        assert output == "这个叫苏向夜."
diff --git a/tests/test_register.py b/tests/test_register.py
index d053e9eb..0ce8351b 100644
--- a/tests/test_register.py
+++ b/tests/test_register.py
@@ -1,5 +1,5 @@
 from infini.core import Core
-from infini.generator import Generator
+from infini.generator import TextGenerator
 from infini.handler import Handler
 from infini.input import Input
 from infini.interceptor import Interceptor
@@ -38,7 +38,7 @@ def test_register():
     pre_interceptor.interceptors = register.pre_interceptors
     handler = Handler()
     handler.handlers = register.handlers
-    generator = Generator()
+    generator = TextGenerator()
     generator.events = register.events
     generator.global_variables = register.global_variables
     interceptor = Interceptor()
-- 
cgit v1.2.3-70-g09d2