1 files changed, 24 insertions, 1 deletions
diff --git a/tests/test_generator.py b/tests/test_generator.py
index 5994e768..ce33f4ff 100644
--- a/tests/test_generator.py
+++ b/tests/test_generator.py
@@ -1,4 +1,4 @@
-from infini.generator import TextGenerator
+from infini.generator import Generator, TextGenerator
 from infini.injector import Injector
 from infini.output import Output
 
@@ -44,3 +44,26 @@ def test_generator_injector():
         )
         == "[苏向夜]Event1 文本: 变量测试"
     )
+
+
+def test_register_generator():
+    def name(nickname: str = "苏向夜"):
+        return nickname
+
+    custom = TextGenerator()
+    custom.type = "custom_text"
+
+    generator = Generator()
+    generator.events = {
+        "test.event1": "[{{ card_name }}]Event1 文本: {{ var }}",
+    }
+    generator.generators.update({"custom_text": custom})
+
+    generator.global_variables = {"card_name": name}
+    assert (
+        generator.output(
+            Output("custom_text", "test.event1", variables={"var": "变量测试"}),
+            Injector(),
+        )
+        == "[苏向夜]Event1 文本: 变量测试"
+    )