Make more use of type inference in tests too

strawberry-graphql · Nov 17, 2024 · 072a299 · 072a299
1 parent b79a106
commit 072a299
Show file tree

Hide file tree

Showing 2 changed files with 274 additions and 343 deletions.
diff --git a/tests/http/clients/base.py b/tests/http/clients/base.py
@@ -24,6 +24,9 @@
 from strawberry.subscriptions.protocols.graphql_transport_ws.handlers import (
     BaseGraphQLTransportWSHandler,
 )
+from strawberry.subscriptions.protocols.graphql_transport_ws.types import (
+    Message as GraphQLTransportWSMessage,
+)
 from strawberry.subscriptions.protocols.graphql_ws.handlers import BaseGraphQLWSHandler
 from strawberry.types import ExecutionResult
 
@@ -301,6 +304,9 @@ async def __aiter__(self) -> AsyncGenerator[Message, None]:
         while not self.closed:
             yield await self.receive()
 
+    async def send_message(self, message: GraphQLTransportWSMessage) -> None:
+        await self.send_json(message)
+
 
 class DebuggableGraphQLTransportWSHandler(BaseGraphQLTransportWSHandler):
     def on_init(self) -> None: