diff --git a/lib/graphql_ws_client.ex b/lib/graphql_ws_client.ex
index 90b036e..f74ac15 100644
--- a/lib/graphql_ws_client.ex
+++ b/lib/graphql_ws_client.ex
@@ -80,8 +80,6 @@ defmodule GraphQLWSClient do
     State
   }
 
-  @client_timeout :infinity
-
   @typedoc """
   Type for a client process.
   """
@@ -181,7 +179,7 @@ defmodule GraphQLWSClient do
   """
   @spec connected?(client) :: boolean
   def connected?(client) do
-    Connection.call(client, :connected?, @client_timeout)
+    call_no_timeout(client, :connected?)
   end
 
   @doc """
@@ -189,7 +187,7 @@ defmodule GraphQLWSClient do
   """
   @spec open(client) :: :ok | {:error, Exception.t()}
   def open(client) do
-    Connection.call(client, :open, @client_timeout)
+    call_no_timeout(client, :open)
   end
 
   @doc """
@@ -197,7 +195,9 @@ defmodule GraphQLWSClient do
   """
   @spec open!(client) :: :ok | no_return
   def open!(client) do
-    client |> open() |> bang!()
+    client
+    |> open()
+    |> bang!()
   end
 
   @doc """
@@ -206,7 +206,7 @@ defmodule GraphQLWSClient do
   @doc since: "1.0.0"
   @spec open_with(client, any) :: :ok | {:error, Exception.t()}
   def open_with(client, init_payload) do
-    Connection.call(client, {:open_with, init_payload}, @client_timeout)
+    call_no_timeout(client, {:open_with, init_payload})
   end
 
   @doc """
@@ -215,7 +215,9 @@ defmodule GraphQLWSClient do
   @doc since: "1.0.0"
   @spec open_with!(client, any) :: :ok | no_return
   def open_with!(client, init_payload) do
-    client |> open_with(init_payload) |> bang!()
+    client
+    |> open_with(init_payload)
+    |> bang!()
   end
 
   @doc """
@@ -223,7 +225,7 @@ defmodule GraphQLWSClient do
   """
   @spec close(client) :: :ok
   def close(client) do
-    Connection.call(client, :close, @client_timeout)
+    call_no_timeout(client, :close)
   end
 
   @doc """
@@ -241,11 +243,7 @@ defmodule GraphQLWSClient do
   @spec query(client, query, variables, nil | timeout) ::
           {:ok, any} | {:error, Exception.t()}
   def query(client, query, variables \\ %{}, timeout \\ nil) do
-    Connection.call(
-      client,
-      {:query, query, variables, timeout},
-      @client_timeout
-    )
+    call_no_timeout(client, {:query, query, variables, timeout})
   end
 
   @doc """
@@ -261,9 +259,9 @@ defmodule GraphQLWSClient do
       ...> )
       %{"data" => %{"posts" => %{"body" => "Lorem Ipsum"}}}
   """
-  @spec query!(client, query, variables, nil | timeout) :: any | no_return
-  def query!(client, query, variables \\ %{}, timeout \\ nil) do
-    case query(client, query, variables, timeout) do
+  @spec query!(client, query, variables) :: any | no_return
+  def query!(client, query, variables \\ %{}) do
+    case query(client, query, variables) do
       {:ok, result} -> result
       {:error, error} -> raise error
     end
@@ -289,11 +287,7 @@ defmodule GraphQLWSClient do
   @spec subscribe(client, query, variables, pid) ::
           {:ok, subscription_id} | {:error, Exception.t()}
   def subscribe(client, query, variables \\ %{}, listener \\ self()) do
-    Connection.call(
-      client,
-      {:subscribe, query, variables, listener},
-      @client_timeout
-    )
+    call_no_timeout(client, {:subscribe, query, variables, listener})
   end
 
   @doc """
@@ -330,7 +324,7 @@ defmodule GraphQLWSClient do
   """
   @spec unsubscribe(client, subscription_id) :: :ok | {:error, Exception.t()}
   def unsubscribe(client, subscription_id) do
-    Connection.call(client, {:unsubscribe, subscription_id}, @client_timeout)
+    call_no_timeout(client, {:unsubscribe, subscription_id})
   end
 
   @doc """
@@ -343,7 +337,9 @@ defmodule GraphQLWSClient do
   """
   @spec unsubscribe!(client, subscription_id) :: :ok | no_return
   def unsubscribe!(client, subscription_id) do
-    client |> unsubscribe(subscription_id) |> bang!()
+    client
+    |> unsubscribe(subscription_id)
+    |> bang!()
   end
 
   @doc """
@@ -378,7 +374,8 @@ defmodule GraphQLWSClient do
       iex> stream |> Stream.take(3) |> Enum.to_list()
   """
   @doc since: "2.0.0"
-  @spec stream!(client, query, variables, Keyword.t()) :: Enumerable.t()
+  @spec stream!(client, query, variables, Keyword.t()) ::
+          Enumerable.t() | no_return
   def stream!(client, query, variables \\ %{}, opts \\ []) do
     Stream.resource(
       fn -> Iterator.open!(client, query, variables, opts) end,
@@ -399,6 +396,10 @@ defmodule GraphQLWSClient do
     }
   end
 
+  defp call_no_timeout(client, req) do
+    Connection.call(client, req, :infinity)
+  end
+
   # Callbacks
 
   @impl true
diff --git a/lib/graphql_ws_client/iterator.ex b/lib/graphql_ws_client/iterator.ex
index a884716..9dc59b7 100644
--- a/lib/graphql_ws_client/iterator.ex
+++ b/lib/graphql_ws_client/iterator.ex
@@ -26,7 +26,11 @@ defmodule GraphQLWSClient.Iterator do
   def open!(client, query, variables \\ %{}, opts \\ []) do
     opts =
       opts
-      |> Keyword.merge(client: client, query: query, variables: variables)
+      |> Keyword.merge(
+        client: client,
+        query: query,
+        variables: variables
+      )
       |> Opts.new()
       |> Opts.validate!()
 
diff --git a/lib/graphql_ws_client/iterator/opts.ex b/lib/graphql_ws_client/iterator/opts.ex
index 81fc6dc..abf3159 100644
--- a/lib/graphql_ws_client/iterator/opts.ex
+++ b/lib/graphql_ws_client/iterator/opts.ex
@@ -3,7 +3,12 @@ defmodule GraphQLWSClient.Iterator.Opts do
 
   @enforce_keys [:client, :query]
 
-  defstruct [:client, :query, buffer_size: 1000, variables: %{}]
+  defstruct [
+    :client,
+    :query,
+    buffer_size: 1000,
+    variables: %{}
+  ]
 
   @type t :: %__MODULE__{
           buffer_size: pos_integer() | :infinity,
diff --git a/test/graphql_ws_client_test.exs b/test/graphql_ws_client_test.exs
index c71a5a4..8e3552a 100644
--- a/test/graphql_ws_client_test.exs
+++ b/test/graphql_ws_client_test.exs
@@ -394,6 +394,8 @@ defmodule GraphQLWSClientTest do
       config = %{@config | query_timeout: 200}
       conn = %{@conn | config: config}
 
+      stub(MockDriver, :disconnect, fn _ -> :ok end)
+
       MockDriver
       |> expect(:connect, fn _ -> {:ok, conn} end)
       |> expect(:push_message, fn _, %Message{id: id} ->
@@ -497,6 +499,8 @@ defmodule GraphQLWSClientTest do
     test "timeout" do
       test_pid = self()
 
+      stub(MockDriver, :disconnect, fn _ -> :ok end)
+
       MockDriver
       |> expect(:connect, fn _ -> {:ok, @conn} end)
       |> expect(:push_message, fn _, %Message{id: id} ->