✅ Add tests for OneBotImplementation

Author: j1g5awi

nonebot_plugin_all4one/middlewares/__init__.py

@@ -2,8 +2,8 @@
 from uuid import uuid4
 from datetime import datetime
 from functools import partial
-from abc import ABC, abstractmethod
 from asyncio import Queue as BaseQueue
+from abc import ABC, ABCMeta, abstractmethod
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -88,8 +88,21 @@ async def get(self):
         return event
 
 
-class Middleware(ABC):
-    supported_actions: ClassVar[Set[str]] = set()
+class _MiddlewareMeta(type):
+    def __new__(cls, name, bases, attrs):
+        supported_actions = set()
+        for base in bases:
+            supported_actions.update(base.supported_actions)
+        attrs["supported_actions"] = supported_actions
+        return type.__new__(cls, name, bases, attrs)
+
+
+class MiddlewareMeta(_MiddlewareMeta, ABCMeta):
+    pass
+
+
+class Middleware(metaclass=MiddlewareMeta):
+    supported_actions: ClassVar[Set[str]]
 
     def __init__(self, bot: Bot):
         self.bot = bot

nonebot_plugin_all4one/onebotimpl/__init__.py

@@ -5,7 +5,18 @@
 from datetime import datetime
 from functools import partial
 from contextlib import asynccontextmanager
-from typing import Any, Dict, List, Type, Union, Literal, Optional, AsyncGenerator, cast
+from typing import (
+    Any,
+    Set,
+    Dict,
+    List,
+    Type,
+    Union,
+    Literal,
+    Optional,
+    AsyncGenerator,
+    cast,
+)
 
 import msgpack
 from nonebot import Driver
@@ -311,7 +322,7 @@ async def _handle_ws(
         await t2
         t1.cancel()
 
-    async def bot_connect(self, bot: Bot) -> None:
+    def bot_connect(self, bot: Bot) -> None:
         if (middleware := self._middlewares.get(bot.type, None)) is None:
             return
         middleware = middleware(bot)
@@ -444,36 +455,32 @@ async def _websocket_rev(
                 )
             await asyncio.sleep(conn.reconnect_interval)
 
-    async def bot_disconnect(self, bot: Bot) -> None:
+    def bot_disconnect(self, bot: Bot) -> None:
         if (middleware := self.middlewares.pop(bot.self_id, None)) is None:
             return
         for task in middleware.tasks:
             if not task.done():
                 task.cancel()
-        await asyncio.gather(*middleware.tasks)
+
+    def import_middlewares(self, middlewares: Optional[Set[str]] = None):
+        if middlewares is None:
+            middlewares = set(self.driver._adapters.keys())
+        for middleware in middlewares:
+            try:
+                if middleware in MIDDLEWARE_MAP:
+                    module = importlib.import_module(
+                        f"nonebot_plugin_all4one.middlewares.{MIDDLEWARE_MAP[middleware]}"
+                    )
+                    self.register_middleware(getattr(module, "Middleware"))
+                else:
+                    logger.warning(f"Can not find middleware for Adapter {middleware}")
+            except Exception as e:
+                logger.warning(f"Can not load middleware for Adapter {middleware}: {e}")
 
     def setup(self):
         @self.driver.on_startup
         async def _():
-            adapters = (
-                self.driver._adapters.keys()
-                if self.config.middlewares is None
-                else self.config.middlewares
-            )
-
-            for adapter in adapters:
-                try:
-                    if adapter in MIDDLEWARE_MAP:
-                        module = importlib.import_module(
-                            f"nonebot_plugin_all4one.middlewares.{MIDDLEWARE_MAP[adapter]}"
-                        )
-                        self.register_middleware(getattr(module, "Middleware"))
-                    else:
-                        logger.warning(f"Can not find middleware for Adapter {adapter}")
-                except Exception as e:
-                    logger.warning(
-                        f"Can not load middleware for Adapter {adapter}: {e}"
-                    )
+            self.import_middlewares(self.config.middlewares)
 
         @self.driver.on_shutdown
         async def _():
@@ -487,10 +494,10 @@ async def _():
         async def _(bot: Bot):
             if bot.self_id.startswith("a4o@"):
                 return
-            await self.bot_connect(bot)
+            self.bot_connect(bot)
 
         @self.driver.on_bot_disconnect
         async def _(bot: Bot):
             if bot.self_id.startswith("a4o@"):
                 return
-            await self.bot_disconnect(bot)
+            self.bot_disconnect(bot)

tests/conftest.py

@@ -27,6 +27,29 @@ def load_adapters(nonebug_init: None):
     driver.register_adapter(TelegramAdapter)
     driver.register_adapter(ConsoleAdapter)
 
+    nonebot.require("nonebot_plugin_all4one")
+    from nonebot_plugin_all4one import obimpl
+
+    obimpl.import_middlewares()
+
+
+@pytest.fixture
+def FakeMiddleware():
+    from nonebot_plugin_all4one.middlewares import Middleware
+
+    class FakeMiddleware(Middleware):
+        @classmethod
+        def get_name(cls):
+            return "fake"
+
+        def get_platform(self):
+            return "fake"
+
+        async def to_onebot_event(self, event):
+            return []
+
+    return FakeMiddleware
+
 
 @pytest.fixture
 async def app(nonebug_init: None, tmp_path: Path):

tests/onebot/events.json tests/middlewares/onebot/events.jsontests/onebot/events.json renamed to tests/middlewares/onebot/events.json

@@ -0,0 +1,13 @@
+from nonebug import App
+from nonebot import get_driver
+
+
+async def test_bot_connect(app: App, FakeMiddleware):
+    from nonebot_plugin_all4one import obimpl
+
+    obimpl.register_middleware(FakeMiddleware)
+
+    async with app.test_api() as ctx:
+        bot = ctx.create_bot()
+        obimpl.bot_connect(bot)
+        assert obimpl.middlewares[bot.self_id].bot == bot