From 8b87e881e5d58d82832333509e626a355647e473 Mon Sep 17 00:00:00 2001
From: Ju4tCode <42488585+yanyongyu@users.noreply.github.com>
Date: Tue, 29 Oct 2024 19:36:44 +0800
Subject: [PATCH 1/7] :bug: fix global asgi test client

---
 nonebug/fixture.py      | 63 ++++++++++++++++++++++++++++-------------
 nonebug/mixin/driver.py | 27 ++++++++++++++++--
 2 files changed, 67 insertions(+), 23 deletions(-)

diff --git a/nonebug/fixture.py b/nonebug/fixture.py
index 006cc4c..a046e8d 100644
--- a/nonebug/fixture.py
+++ b/nonebug/fixture.py
@@ -1,10 +1,50 @@
+from contextlib import asynccontextmanager, nullcontext
+
 import pytest
+from async_asgi_testclient import TestClient
 
 from nonebug.app import App
+from nonebug.mixin.driver import set_global_client
 
 from . import NONEBOT_INIT_KWARGS, NONEBOT_START_LIFESPAN
 
 
+@asynccontextmanager
+async def lifespan_ctx():
+    import nonebot
+    from nonebot.drivers import ASGIMixin
+
+    driver = nonebot.get_driver()
+
+    if isinstance(driver, ASGIMixin):
+        # if the driver has an asgi application
+        # use asgi lifespan to startup/shutdown
+        ctx = TestClient(driver.asgi)
+        set_global_client(ctx)
+    else:
+        ctx = driver._lifespan
+
+    try:
+        await ctx.__aenter__()
+    except Exception as e:
+        logger.opt(colors=True, exception=e).error(
+            "<r><bg #f8bbd0>Error occurred while running startup hook."
+            "</bg #f8bbd0></r>"
+        )
+        raise
+
+    try:
+        yield
+    finally:
+        try:
+            await ctx.__aexit__(None, None, None)
+        except Exception as e:
+            logger.opt(colors=True, exception=e).error(
+                "<r><bg #f8bbd0>Error occurred while running shutdown hook."
+                "</bg #f8bbd0></r>"
+            )
+
+
 @pytest.fixture(scope="session", autouse=True)
 async def nonebug_init(request: pytest.FixtureRequest):  # noqa: PT004
     """
@@ -20,28 +60,11 @@ async def nonebug_init(request: pytest.FixtureRequest):  # noqa: PT004
     matchers.set_provider(NoneBugProvider)
 
     run_lifespan = request.config.stash.get(NONEBOT_START_LIFESPAN, True)
-    driver = nonebot.get_driver()
-    if run_lifespan:
-        try:
-            await driver._lifespan.startup()
-        except Exception as e:
-            logger.opt(colors=True, exception=e).error(
-                "<r><bg #f8bbd0>Error occurred while running startup hook."
-                "</bg #f8bbd0></r>"
-            )
-            raise
 
-    try:
+    ctx = lifespan_ctx() if run_lifespan else nullcontext()
+
+    async with ctx:
         yield
-    finally:
-        if run_lifespan:
-            try:
-                await driver._lifespan.shutdown()
-            except Exception as e:
-                logger.opt(colors=True, exception=e).error(
-                    "<r><bg #f8bbd0>Error occurred while running shutdown hook."
-                    "</bg #f8bbd0></r>"
-                )
 
 
 @pytest.fixture(name="app")
diff --git a/nonebug/mixin/driver.py b/nonebug/mixin/driver.py
index c14b1f8..eedd6d4 100644
--- a/nonebug/mixin/driver.py
+++ b/nonebug/mixin/driver.py
@@ -7,6 +7,20 @@
 from nonebug.base import BaseApp, Context
 
 
+_global_client: Optional[TestClient] = None
+
+
+def set_global_client(client: TestClient):
+    global _global_client
+
+    if _global_client is not None:
+        raise RuntimeError()
+
+
+def get_global_client() -> Optional[TestClient]:
+    return _global_client
+
+
 @final
 class ServerContext(Context):
     def __init__(
@@ -14,18 +28,21 @@ def __init__(
         app: BaseApp,
         *args,
         asgi: ASGIApplication,
+        client: Optional[TestClient] = None,
         **kwargs,
     ):
         super().__init__(app, *args, **kwargs)
         self.asgi = asgi
+        self.specified_client = client
         self.client = TestClient(self.asgi)
 
     def get_client(self) -> TestClient:
-        return self.client
+        return self.specified_client or self.client
 
     async def setup(self) -> None:
         await super().setup()
-        await self.stack.enter_async_context(self.client)
+        if self.specified_client is None:
+            await self.stack.enter_async_context(self.client)
 
 
 # @final
@@ -50,8 +67,12 @@ class DriverMixin(BaseApp):
     def test_server(self, asgi: Optional[ASGIApplication] = None) -> ServerContext:
         import nonebot
 
+        client = None
+        if asgi is None:
+            client = get_global_client()
+
         asgi = asgi or nonebot.get_asgi()
-        return ServerContext(self, asgi=asgi)
+        return ServerContext(self, asgi=asgi, client=client)
 
     # def test_client(self):
     #     ...

From b493a2862673b94a17b0266ac8ea114d84ad52ab Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Tue, 29 Oct 2024 11:38:54 +0000
Subject: [PATCH 2/7] :rotating_light: auto fix by pre-commit hooks

---
 nonebug/fixture.py      | 3 +--
 nonebug/mixin/driver.py | 1 -
 2 files changed, 1 insertion(+), 3 deletions(-)

diff --git a/nonebug/fixture.py b/nonebug/fixture.py
index a046e8d..3e6c813 100644
--- a/nonebug/fixture.py
+++ b/nonebug/fixture.py
@@ -1,4 +1,4 @@
-from contextlib import asynccontextmanager, nullcontext
+from contextlib import nullcontext, asynccontextmanager
 
 import pytest
 from async_asgi_testclient import TestClient
@@ -51,7 +51,6 @@ async def nonebug_init(request: pytest.FixtureRequest):  # noqa: PT004
     Initialize nonebot before test case running.
     """
     import nonebot
-    from nonebot import logger
     from nonebot.matcher import matchers
 
     from nonebug.provider import NoneBugProvider
diff --git a/nonebug/mixin/driver.py b/nonebug/mixin/driver.py
index eedd6d4..18f5617 100644
--- a/nonebug/mixin/driver.py
+++ b/nonebug/mixin/driver.py
@@ -6,7 +6,6 @@
 
 from nonebug.base import BaseApp, Context
 
-
 _global_client: Optional[TestClient] = None
 
 

From aa2fb4a94ac698cdc10dc03e46fc79d9694a431c Mon Sep 17 00:00:00 2001
From: Ju4tCode <42488585+yanyongyu@users.noreply.github.com>
Date: Tue, 29 Oct 2024 19:40:32 +0800
Subject: [PATCH 3/7] :bug: fix import

---
 nonebug/fixture.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/nonebug/fixture.py b/nonebug/fixture.py
index 3e6c813..f93a32b 100644
--- a/nonebug/fixture.py
+++ b/nonebug/fixture.py
@@ -12,6 +12,7 @@
 @asynccontextmanager
 async def lifespan_ctx():
     import nonebot
+    from nonebot import logger
     from nonebot.drivers import ASGIMixin
 
     driver = nonebot.get_driver()

From 7b757b76d27c162b874d6e442ca2e62350461dad Mon Sep 17 00:00:00 2001
From: Ju4tCode <42488585+yanyongyu@users.noreply.github.com>
Date: Wed, 30 Oct 2024 14:04:51 +0800
Subject: [PATCH 4/7] Update driver.py

---
 nonebug/mixin/driver.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/nonebug/mixin/driver.py b/nonebug/mixin/driver.py
index 18f5617..bc79c77 100644
--- a/nonebug/mixin/driver.py
+++ b/nonebug/mixin/driver.py
@@ -15,6 +15,8 @@ def set_global_client(client: TestClient):
     if _global_client is not None:
         raise RuntimeError()
 
+    _global_client = client
+
 
 def get_global_client() -> Optional[TestClient]:
     return _global_client

From f7dca9ed4c7d5e10146373cf84a063cf7abb94de Mon Sep 17 00:00:00 2001
From: Ju4tCode <42488585+yanyongyu@users.noreply.github.com>
Date: Wed, 30 Oct 2024 06:53:11 +0000
Subject: [PATCH 5/7] :bug: fix startup before init setup

---
 nonebug/fixture.py | 10 +++++++++-
 1 file changed, 9 insertions(+), 1 deletion(-)

diff --git a/nonebug/fixture.py b/nonebug/fixture.py
index f93a32b..90a3acf 100644
--- a/nonebug/fixture.py
+++ b/nonebug/fixture.py
@@ -47,7 +47,7 @@ async def lifespan_ctx():
 
 
 @pytest.fixture(scope="session", autouse=True)
-async def nonebug_init(request: pytest.FixtureRequest):  # noqa: PT004
+def _nonebot_init(request: pytest.FixtureRequest):
     """
     Initialize nonebot before test case running.
     """
@@ -59,6 +59,14 @@ async def nonebug_init(request: pytest.FixtureRequest):  # noqa: PT004
     nonebot.init(**request.config.stash.get(NONEBOT_INIT_KWARGS, {}))
     matchers.set_provider(NoneBugProvider)
 
+
+@pytest.fixture(scope="session", autouse=True)
+async def after_nonebot_init(_nonebot_init: None):  # noqa: PT004
+    pass
+
+
+@pytest.fixture(scope="session", autouse=True)
+async def nonebug_init(after_nonebot_init: None, request: pytest.FixtureRequest):  # noqa: PT004
     run_lifespan = request.config.stash.get(NONEBOT_START_LIFESPAN, True)
 
     ctx = lifespan_ctx() if run_lifespan else nullcontext()

From cb4f075c065e74c7ff545e1a2ff11697e9d97c66 Mon Sep 17 00:00:00 2001
From: Ju4tCode <42488585+yanyongyu@users.noreply.github.com>
Date: Wed, 30 Oct 2024 07:16:41 +0000
Subject: [PATCH 6/7] :bug: fix test

---
 nonebug/fixture.py | 4 +++-
 tests/conftest.py  | 4 ++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/nonebug/fixture.py b/nonebug/fixture.py
index 90a3acf..6eea0ba 100644
--- a/nonebug/fixture.py
+++ b/nonebug/fixture.py
@@ -66,7 +66,9 @@ async def after_nonebot_init(_nonebot_init: None):  # noqa: PT004
 
 
 @pytest.fixture(scope="session", autouse=True)
-async def nonebug_init(after_nonebot_init: None, request: pytest.FixtureRequest):  # noqa: PT004
+async def nonebug_init(  # noqa: PT004
+    _nonebot_init: None, after_nonebot_init: None, request: pytest.FixtureRequest
+):
     run_lifespan = request.config.stash.get(NONEBOT_START_LIFESPAN, True)
 
     ctx = lifespan_ctx() if run_lifespan else nullcontext()
diff --git a/tests/conftest.py b/tests/conftest.py
index 6685658..e3328dd 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -12,6 +12,6 @@ def pytest_configure(config: pytest.Config) -> None:
     config.stash[NONEBOT_INIT_KWARGS] = {"custom_key": "custom_value"}
 
 
-@pytest.fixture(scope="session")
-def load_plugin(nonebug_init: None) -> set[Plugin]:
+@pytest.fixture(scope="session", autouse=True)
+async def after_nonebot_init(_nonebot_init: None) -> set[Plugin]:
     return nonebot.load_plugins(str(Path(__file__).parent / "plugins"))

From 11642e54c43a487ed72c5067d7a9b17ac2c24f2d Mon Sep 17 00:00:00 2001
From: Ju4tCode <42488585+yanyongyu@users.noreply.github.com>
Date: Wed, 30 Oct 2024 08:19:57 +0000
Subject: [PATCH 7/7] :white_check_mark: fix test

---
 nonebug/fixture.py | 2 +-
 tests/conftest.py  | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/nonebug/fixture.py b/nonebug/fixture.py
index 6eea0ba..d7254f7 100644
--- a/nonebug/fixture.py
+++ b/nonebug/fixture.py
@@ -86,4 +86,4 @@ def nonebug_app(nonebug_init) -> App:
     return App()
 
 
-__all__ = ["nonebug_init", "nonebug_app"]
+__all__ = ["after_nonebot_init", "nonebug_init", "nonebug_app"]
diff --git a/tests/conftest.py b/tests/conftest.py
index e3328dd..76ba821 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -5,7 +5,7 @@
 from nonebot.plugin import Plugin
 
 from nonebug import NONEBOT_INIT_KWARGS
-from nonebug.fixture import *  # noqa: F403
+from nonebug.fixture import nonebug_app, nonebug_init, _nonebot_init  # noqa: F401
 
 
 def pytest_configure(config: pytest.Config) -> None:
@@ -13,5 +13,5 @@ def pytest_configure(config: pytest.Config) -> None:
 
 
 @pytest.fixture(scope="session", autouse=True)
-async def after_nonebot_init(_nonebot_init: None) -> set[Plugin]:
+async def after_nonebot_init(_nonebot_init: None) -> set[Plugin]:  # noqa: F811
     return nonebot.load_plugins(str(Path(__file__).parent / "plugins"))