From c4d0b0555cc98ffbb50ea7bd95f40fdcae22b5dd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Kristj=C3=A1n=20Valur=20J=C3=B3nsson?= <sweskman@gmail.com>
Date: Sun, 13 Oct 2024 19:39:47 +0000
Subject: [PATCH] cleanup

---
 RELEASE.md                                    |  3 +--
 tests/http/clients/litestar.py                |  4 ++++
 tests/views/schema.py                         | 14 ++++++-------
 tests/websockets/test_graphql_transport_ws.py | 20 ++++---------------
 4 files changed, 16 insertions(+), 25 deletions(-)

diff --git a/RELEASE.md b/RELEASE.md
index 2820214c74..8f6812b821 100644
--- a/RELEASE.md
+++ b/RELEASE.md
@@ -1,4 +1,3 @@
 Release type: patch
 
-Operations over `graphql-transport-ws` now create the Context and perform validation on
-the worker `Task`, thus not blocking the websocket from accepting messages.
+Fix error handling for query operations over graphql-transport-ws
\ No newline at end of file
diff --git a/tests/http/clients/litestar.py b/tests/http/clients/litestar.py
index b29be52b32..2548dc563c 100644
--- a/tests/http/clients/litestar.py
+++ b/tests/http/clients/litestar.py
@@ -29,6 +29,10 @@
 )
 
 
+def custom_context_dependency() -> str:
+    return "Hi!"
+
+
 async def litestar_get_context(request: Request = None):
     return get_context({"request": request})
 
diff --git a/tests/views/schema.py b/tests/views/schema.py
index d8a4e89107..cf50dfe415 100644
--- a/tests/views/schema.py
+++ b/tests/views/schema.py
@@ -77,7 +77,7 @@ class DebugInfo:
 @strawberry.type
 class Query:
     @strawberry.field
-    def greetings(self) -> str:  # pragma: no cover
+    def greetings(self) -> str:
         return "hello"
 
     @strawberry.field
@@ -91,13 +91,13 @@ async def async_hello(self, name: Optional[str] = None, delay: float = 0) -> str
 
     @strawberry.field(permission_classes=[AlwaysFailPermission])
     def always_fail(self) -> Optional[str]:
-        return "Hey"  # pragma: no cover
+        return "Hey"
 
     @strawberry.field(permission_classes=[ConditionalFailPermission])
     def conditional_fail(
         self, sleep: Optional[float] = None, fail: bool = False
     ) -> str:
-        return "Hey"  # pragma: no cover
+        return "Hey"
 
     @strawberry.field
     async def error(self, message: str) -> AsyncGenerator[str, None]:
@@ -108,7 +108,7 @@ async def exception(self, message: str) -> str:
         raise ValueError(message)
 
     @strawberry.field
-    def teapot(self, info: strawberry.Info[Any, None]) -> str:  # pragma: no cover
+    def teapot(self, info: strawberry.Info[Any, None]) -> str:
         info.context["response"].status_code = 418
 
         return "🫖"
@@ -142,7 +142,7 @@ def set_header(self, info: strawberry.Info, name: str) -> str:
 @strawberry.type
 class Mutation:
     @strawberry.mutation
-    def echo(self, string_to_echo: str) -> str:  # pragma: no cover
+    def echo(self, string_to_echo: str) -> str:
         return string_to_echo
 
     @strawberry.mutation
@@ -162,7 +162,7 @@ def read_folder(self, folder: FolderInput) -> List[str]:
         return list(map(_read_file, folder.files))
 
     @strawberry.mutation
-    def match_text(self, text_file: Upload, pattern: str) -> str:  # pragma: no cover
+    def match_text(self, text_file: Upload, pattern: str) -> str:
         text = text_file.read().decode()
         return pattern if pattern in text else ""
 
@@ -199,7 +199,7 @@ async def exception(self, message: str) -> AsyncGenerator[str, None]:
         raise ValueError(message)
 
         # Without this yield, the method is not recognised as an async generator
-        yield "Hi"  # pragma: no cover
+        yield "Hi"
 
     @strawberry.subscription
     async def flavors(self) -> AsyncGenerator[Flavor, None]:
diff --git a/tests/websockets/test_graphql_transport_ws.py b/tests/websockets/test_graphql_transport_ws.py
index e4cd0ad109..fb4541978c 100644
--- a/tests/websockets/test_graphql_transport_ws.py
+++ b/tests/websockets/test_graphql_transport_ws.py
@@ -30,23 +30,8 @@
 from tests.http.clients.base import DebuggableGraphQLTransportWSHandler
 from tests.views.schema import MyExtension, Schema
 
-from ..http.clients.base import WebSocketClient
-
-try:
-    from ..http.clients.fastapi import FastAPIHttpClient
-except ImportError:  # pragma: no cover
-    FastAPIHttpClient = None
-try:
-    from ..http.clients.starlite import StarliteHttpClient
-except ImportError:  # pragma: no cover
-    StarliteHttpClient = None
-try:
-    from ..http.clients.litestar import LitestarHttpClient
-except ImportError:  # pragma: no cover
-    LitestarHttpClient = None
-
 if TYPE_CHECKING:
-    from ..http.clients.base import HttpClient
+    from ..http.clients.base import HttpClient, WebSocketClient
 
 
 @pytest_asyncio.fixture
@@ -913,6 +898,9 @@ async def test_error_handler_for_timeout(http_client: HttpClient):
         if isinstance(http_client, ChannelsHttpClient):
             pytest.skip("Can't patch on_init for this client")
 
+    if not AsyncMock:
+        pytest.skip("Don't have AsyncMock")
+
     ws = ws_raw
     handler = None
     errorhandler = AsyncMock()