From b037be448537be9d430078289780cd1183fbc4c2 Mon Sep 17 00:00:00 2001
From: Mix <32300164+mnixry@users.noreply.github.com>
Date: Mon, 17 Jan 2022 00:29:09 +0800
Subject: [PATCH] :white_check_mark: add unit test for message slice

---
 tests/test_adapters/test_message.py  | 54 ++++++++++++++++++++++++++++
 tests/test_adapters/test_template.py | 17 ---------
 tests/utils.py                       | 11 ++++--
 3 files changed, 62 insertions(+), 20 deletions(-)
 create mode 100644 tests/test_adapters/test_message.py
 delete mode 100644 tests/test_adapters/test_template.py

diff --git a/tests/test_adapters/test_message.py b/tests/test_adapters/test_message.py
new file mode 100644
index 00000000..88d9e563
--- /dev/null
+++ b/tests/test_adapters/test_message.py
@@ -0,0 +1,54 @@
+from utils import make_fake_message
+
+
+def test_message_template():
+    from nonebot.adapters import MessageTemplate
+
+    Message = make_fake_message()
+
+    template = MessageTemplate("{a:custom}{b:text}{c:image}", Message)
+
+    @template.add_format_spec
+    def custom(input: str) -> str:
+        return input + "-custom!"
+
+    formatted = template.format(a="test", b="test", c="https://example.com/test")
+    assert formatted.extract_plain_text() == "test-custom!test"
+    assert str(formatted) == "test-custom!test[fake:image]"
+
+
+def test_message_slice():
+
+    Message = make_fake_message()
+    MessageSegment = Message.get_segment_class()
+
+    message = Message(
+        [
+            MessageSegment.text("test"),
+            MessageSegment.image("test2"),
+            MessageSegment.image("test3"),
+            MessageSegment.text("test4"),
+        ]
+    )
+
+    assert message[0] == MessageSegment.text("test")
+
+    assert message[0:2] == Message(
+        [MessageSegment.text("test"), MessageSegment.image("test2")]
+    )
+
+    assert message["image"] == Message(
+        [MessageSegment.image("test2"), MessageSegment.image("test3")]
+    )
+
+    assert message["image", 0] == MessageSegment.image("test2")
+    assert message["image", 0:2] == message["image"]
+
+    assert message.index(message[0]) == 0
+    assert message.index("image") == 1
+
+    assert message.get("image") == message["image"]
+    assert message.get("image", 114514) == message["image"]
+    assert message.get("image", 1) == Message([message["image", 0]])
+
+    assert message.count("image") == 2
diff --git a/tests/test_adapters/test_template.py b/tests/test_adapters/test_template.py
deleted file mode 100644
index 3dbef541..00000000
--- a/tests/test_adapters/test_template.py
+++ /dev/null
@@ -1,17 +0,0 @@
-from utils import make_fake_message
-
-
-def test_message_template():
-    from nonebot.adapters import MessageTemplate
-
-    Message = make_fake_message()
-
-    template = MessageTemplate("{a:custom}{b:text}{c:image}", Message)
-
-    @template.add_format_spec
-    def custom(input: str) -> str:
-        return input + "-custom!"
-
-    formatted = template.format(a="test", b="test", c="https://example.com/test")
-    assert formatted.extract_plain_text() == "test-custom!test"
-    assert str(formatted) == "test-custom!test[fake:image]"
diff --git a/tests/utils.py b/tests/utils.py
index 09b2a987..711c7c05 100644
--- a/tests/utils.py
+++ b/tests/utils.py
@@ -1,4 +1,4 @@
-from typing import TYPE_CHECKING, Type, Optional
+from typing import TYPE_CHECKING, Type, Union, Mapping, Iterable, Optional
 
 from pydantic import create_model
 
@@ -34,8 +34,13 @@ def make_fake_message() -> Type["Message"]:
             return FakeMessageSegment
 
         @staticmethod
-        def _construct(msg: str):
-            yield FakeMessageSegment.text(msg)
+        def _construct(msg: Union[str, Iterable[Mapping]]):
+            if isinstance(msg, str):
+                yield FakeMessageSegment.text(msg)
+            else:
+                for seg in msg:
+                    yield FakeMessageSegment(**seg)
+            return
 
     return FakeMessage