Added TestClient.get_stats

mongkok · mongkok · commit 180eab8bbcc8 · 2021-08-24T06:10:48.000+07:00
diff --git a/tests/panels/sqlalchemy/test_sqlalchemy.py b/tests/panels/sqlalchemy/test_sqlalchemy.py
@@ -1,5 +1,3 @@
-from fastapi import status
-
 from ...mark import override_panels, skip_py
 from ...testclient import TestClient
 
@@ -8,12 +6,9 @@
 @override_panels(["debug_toolbar.panels.sqlalchemy.SQLAlchemyPanel"])
 def test_sqlalchemy(client: TestClient) -> None:
     store_id = client.get_store_id("/sql")
-    response = client.render_panel(store_id, "SQLAlchemyPanel")
-
-    assert response.status_code == status.HTTP_200_OK
+    stats = client.get_stats(store_id, "SQLAlchemyPanel")
+    queries = stats["queries"]
 
-    content = response.json()["content"]
-    assert "3 queries" in content
-    assert "2 similar" in content
-    assert "INSERT" in content
-    assert "SELECT" in content
+    assert len(queries) == 3
+    assert queries[0][1]["sql"].startswith("INSERT")
+    assert queries[1][1]["dup_count"] == 2
diff --git a/tests/panels/test_headers.py b/tests/panels/test_headers.py
@@ -1,13 +1,14 @@
-from fastapi import status
-
 from ..mark import override_panels
 from ..testclient import TestClient
 
 
 @override_panels(["debug_toolbar.panels.headers.HeadersPanel"])
 def test_headers(client: TestClient) -> None:
-    store_id = client.get_store_id("/async")
-    response = client.render_panel(store_id, "HeadersPanel")
-
-    assert response.status_code == status.HTTP_200_OK
-    assert "accept" in response.json()["content"]
+    headers = {
+        "cookie": "",
+    }
+    store_id = client.get_store_id("/async", headers=headers)
+    stats = client.get_stats(store_id, "HeadersPanel")
+    request_headers = stats["request_headers"]
+
+    assert request_headers["cookie"]
diff --git a/tests/panels/test_logging.py b/tests/panels/test_logging.py
@@ -2,7 +2,7 @@
 import typing as t
 
 import pytest
-from fastapi import FastAPI, Request, status
+from fastapi import FastAPI, Request
 from fastapi.logger import logger
 from fastapi.responses import HTMLResponse
 
@@ -29,7 +29,6 @@ async def get_log_async(request: Request, level: str) -> str:
 @override_panels(["debug_toolbar.panels.logging.LoggingPanel"])
 def test_logging(client: TestClient, path: str, level: str) -> None:
     store_id = client.get_store_id(f"/log/{path}?level={level}")
-    response = client.render_panel(store_id, "LoggingPanel")
+    stats = client.get_stats(store_id, "LoggingPanel")
 
-    assert response.status_code == status.HTTP_200_OK
-    assert level in response.json()["content"]
+    assert stats["records"][0]["level"] == level
diff --git a/tests/panels/test_profiling.py b/tests/panels/test_profiling.py
@@ -1,5 +1,4 @@
 import pytest
-from fastapi import status
 
 from ..mark import override_panels
 from ..testclient import TestClient
@@ -9,7 +8,6 @@
 @override_panels(["debug_toolbar.panels.profiling.ProfilingPanel"])
 def test_profiling(client: TestClient, path: str) -> None:
     store_id = client.get_store_id(f"/{path}")
-    response = client.render_panel(store_id, "ProfilingPanel")
+    stats = client.get_stats(store_id, "ProfilingPanel")
 
-    assert response.status_code == status.HTTP_200_OK
-    assert "profileSession" in response.json()["content"]
+    assert "profileSession" in stats["content"]
diff --git a/tests/panels/test_pydantic.py b/tests/panels/test_pydantic.py
@@ -1,31 +1,10 @@
-import typing as t
-
-import pytest
-from fastapi import Depends, FastAPI, Request, status
-from fastapi.responses import HTMLResponse
-from pydantic import BaseModel
-
 from ..mark import override_panels
 from ..testclient import TestClient
 
 
-class Model(BaseModel):
-    test: bool
-
-
-@pytest.fixture
-def client(app: FastAPI, get_index: t.Callable) -> TestClient:
-    @app.get("/pydantic", response_class=HTMLResponse)
-    async def get_pydantic(request: Request, model: Model = Depends()) -> str:
-        return get_index(request)
-
-    return TestClient(app)
-
-
 @override_panels(["debug_toolbar.panels.pydantic.PydanticPanel"])
 def test_pydantic(client: TestClient) -> None:
-    store_id = client.get_store_id("/pydantic", params={"test": True})
-    response = client.render_panel(store_id, "PydanticPanel")
+    store_id = client.get_store_id("/openapi.json")
+    stats = client.get_stats(store_id, "PydanticPanel")
 
-    assert response.status_code == status.HTTP_200_OK
-    assert "Model.test" in response.json()["content"]
+    assert stats["validations"]
diff --git a/tests/panels/test_request.py b/tests/panels/test_request.py
@@ -1,7 +1,7 @@
 import typing as t
 
 import pytest
-from fastapi import FastAPI, Request, status
+from fastapi import FastAPI, Request
 from fastapi.responses import HTMLResponse
 from starlette.middleware.sessions import SessionMiddleware
 
@@ -24,7 +24,6 @@ async def get_session(request: Request) -> str:
 @override_panels(["debug_toolbar.panels.request.RequestPanel"])
 def test_session(client: TestClient) -> None:
     store_id = client.get_store_id("/session")
-    response = client.render_panel(store_id, "RequestPanel")
+    stats = client.get_stats(store_id, "RequestPanel")
 
-    assert response.status_code == status.HTTP_200_OK
-    assert "debug" in response.json()["content"]
+    assert stats["session"]["debug"]
diff --git a/tests/panels/test_routes.py b/tests/panels/test_routes.py
@@ -1,13 +1,10 @@
-from fastapi import status
-
 from ..mark import override_panels
 from ..testclient import TestClient
 
 
 @override_panels(["debug_toolbar.panels.routes.RoutesPanel"])
 def test_routes(client: TestClient) -> None:
     store_id = client.get_store_id("/async")
-    response = client.render_panel(store_id, "RoutesPanel")
+    stats = client.get_stats(store_id, "RoutesPanel")
 
-    assert response.status_code == status.HTTP_200_OK
-    assert "openapi" in response.json()["content"]
+    assert any(route.name == "openapi" for route in stats["routes"])
diff --git a/tests/panels/test_settings.py b/tests/panels/test_settings.py
@@ -1,4 +1,4 @@
-from fastapi import status
+import pytest
 
 from ..mark import override_panels
 from ..testclient import TestClient
@@ -7,7 +7,6 @@
 @override_panels(["debug_toolbar.panels.settings.SettingsPanel"])
 def test_settings(client: TestClient) -> None:
     store_id = client.get_store_id("/async")
-    response = client.render_panel(store_id, "SettingsPanel")
 
-    assert response.status_code == status.HTTP_200_OK
-    assert "PANELS" in response.json()["content"]
+    with pytest.raises(KeyError):
+        client.get_stats(store_id, "SettingsPanel")
diff --git a/tests/panels/test_timer.py b/tests/panels/test_timer.py
@@ -1,15 +1,10 @@
-import re
-
-from fastapi import status
-
 from ..mark import override_panels
 from ..testclient import TestClient
 
 
 @override_panels(["debug_toolbar.panels.timer.TimerPanel"])
 def test_timer(client: TestClient) -> None:
     store_id = client.get_store_id("/async")
-    response = client.render_panel(store_id, "TimerPanel")
+    stats = client.get_stats(store_id, "TimerPanel")
 
-    assert response.status_code == status.HTTP_200_OK
-    assert re.findall(r"(\d+) msec", response.json()["content"])
+    assert "total" in stats
diff --git a/tests/panels/test_versions.py b/tests/panels/test_versions.py
@@ -1,13 +1,10 @@
-from fastapi import status
-
 from ..mark import override_panels
 from ..testclient import TestClient
 
 
 @override_panels(["debug_toolbar.panels.versions.VersionsPanel"])
 def test_versions(client: TestClient) -> None:
     store_id = client.get_store_id("/async")
-    response = client.render_panel(store_id, "VersionsPanel")
+    stats = client.get_stats(store_id, "VersionsPanel")
 
-    assert response.status_code == status.HTTP_200_OK
-    assert "fastapi" in response.json()["content"]
+    assert stats["packages"]
diff --git a/tests/panels/tortoise/test_tortoise.py b/tests/panels/tortoise/test_tortoise.py
@@ -1,18 +1,13 @@
-from fastapi import status
-
 from ...mark import override_panels
 from ...testclient import TestClient
 
 
 @override_panels(["debug_toolbar.panels.tortoise.TortoisePanel"])
 def test_tortoise(client: TestClient) -> None:
     store_id = client.get_store_id("/sql")
-    response = client.render_panel(store_id, "TortoisePanel")
-
-    assert response.status_code == status.HTTP_200_OK
+    stats = client.get_stats(store_id, "TortoisePanel")
+    queries = stats["queries"]
 
-    content = response.json()["content"]
-    assert "3 queries" in content
-    assert "2 similar" in content
-    assert "INSERT" in content
-    assert "SELECT" in content
+    assert len(queries) == 3
+    assert queries[0][1]["sql"].startswith("INSERT")
+    assert queries[1][1]["dup_count"] == 2
diff --git a/tests/testclient.py b/tests/testclient.py
@@ -1,15 +1,33 @@
 import re
 import typing as t
+from urllib import parse
 
 import requests
+from fastapi import status
 from fastapi.testclient import TestClient as BaseTestClient
 
+from debug_toolbar.toolbar import DebugToolbar
+from debug_toolbar.types import Stats
+
 
 class TestClient(BaseTestClient):
     def get_store_id(self, path: str, method: str = None, **kwargs: t.Any) -> str:
         response = getattr(self, (method or "get"))(path, **kwargs)
-        return re.findall(r'store-id="(.+?)"', response.content.decode())[0]
+
+        if response.headers["content-type"].startswith("text/html"):
+            return re.findall(r'store-id="(.+?)"', response.content.decode())[0]
+
+        cookie = parse.unquote(response.headers["set-cookie"])
+        return re.findall(r'"storeId": "(.+?)"', cookie)[0]
 
     def render_panel(self, store_id: str, panel_id: str) -> requests.Response:
-        url = self.app.router.url_path_for("debug_toolbar.render_panel")  # type: ignore
+        url = self.app.url_path_for("debug_toolbar.render_panel")  # type: ignore
         return self.get(f"{url}?store_id={store_id}&panel_id={panel_id}")
+
+    def get_stats(self, store_id: str, panel_id: str) -> Stats:
+        response = self.render_panel(store_id, panel_id)
+        assert response.status_code == status.HTTP_200_OK
+
+        toolbar = DebugToolbar.fetch(store_id)
+        assert toolbar
+        return toolbar.stats[panel_id]
