From a2b2114e97b14995b91a74ea74b26008eeeeb263 Mon Sep 17 00:00:00 2001
From: Synrom <max.leiwig@gmail.com>
Date: Sun, 8 Sep 2024 13:11:12 +0200
Subject: [PATCH] Add more tests

---
 tests/test_ignore_trailing_slash.py | 55 ++++++++++++++++++++++++++++-
 1 file changed, 54 insertions(+), 1 deletion(-)

diff --git a/tests/test_ignore_trailing_slash.py b/tests/test_ignore_trailing_slash.py
index 61ef7b496..10830077c 100644
--- a/tests/test_ignore_trailing_slash.py
+++ b/tests/test_ignore_trailing_slash.py
@@ -1,4 +1,5 @@
-from fastapi import APIRouter, FastAPI, WebSocket
+from fastapi import APIRouter, FastAPI, Request, WebSocket
+from fastapi.responses import JSONResponse
 from fastapi.testclient import TestClient
 
 app = FastAPI(ignore_trailing_slash=True)
@@ -53,6 +54,44 @@ def route_endpoint_with_slash():
     return {"msg": "Routing Example 2"}
 
 
+@router.websocket("/websocket")
+async def router_websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    await websocket.send_text("Websocket")
+    await websocket.close()
+
+
+@router.websocket("/websocket2/")
+async def router_websocket_endpoint_with_slash(websocket: WebSocket):
+    await websocket.accept()
+    await websocket.send_text("Websocket 2")
+    await websocket.close()
+
+
+@router.websocket_route("/websocket_route")
+async def router_websocket_route_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    await websocket.send_text("Websocket route")
+    await websocket.close()
+
+
+@router.websocket_route("/websocket_route_2/")
+async def router_websocket_route_endpoint_with_slash(websocket: WebSocket):
+    await websocket.accept()
+    await websocket.send_text("Websocket route 2")
+    await websocket.close()
+
+
+@router.route("/starlette_route", ["get"])
+async def starlette_route_endpoint(request: Request):
+    return JSONResponse({"msg": "Starlette Route"})
+
+
+@router.route("/starlette_route_2/", ["get"])
+async def starlette_route_endpoint_with_slash(request: Request):
+    return JSONResponse({"msg": "Starlette Route 2"})
+
+
 app.include_router(router, prefix="/router")
 
 client = TestClient(app)
@@ -85,3 +124,17 @@ def test_ignoring_trailing_routing():
     response = client.get("router/example2", follow_redirects=False)
     assert response.status_code == 200
     assert response.json()["msg"] == "Routing Example 2"
+    response = client.get("router/starlette_route/", follow_redirects=False)
+    assert response.status_code == 200
+    assert response.json()["msg"] == "Starlette Route"
+    response = client.get("router/starlette_route_2", follow_redirects=False)
+    assert response.status_code == 200
+    assert response.json()["msg"] == "Starlette Route 2"
+    with client.websocket_connect("router/websocket/") as websocket:
+        assert websocket.receive_text() == "Websocket"
+    with client.websocket_connect("router/websocket2") as websocket:
+        assert websocket.receive_text() == "Websocket 2"
+    with client.websocket_connect("router/websocket_route/") as websocket:
+        assert websocket.receive_text() == "Websocket route"
+    with client.websocket_connect("router/websocket_route_2/") as websocket:
+        assert websocket.receive_text() == "Websocket route 2"