fix: trap exits in iterator server

lib/graphql_ws_client/iterator.ex

@@ -62,15 +62,15 @@ defmodule GraphQLWSClient.Iterator do
 
   @impl true
   def init(%Opts{} = opts) do
+    Process.flag(:trap_exit, true)
+
     case GraphQLWSClient.subscribe(opts.client, opts.query, opts.variables) do
       {:ok, subscription_id} ->
-        monitor_ref = Process.monitor(opts.client)
-
         {:ok,
          %State{
            buffer_size: opts.buffer_size,
            client: opts.client,
-           monitor_ref: monitor_ref,
+           monitor_ref: Process.monitor(opts.client),
            subscription_id: subscription_id
          }}
 
@@ -81,7 +81,9 @@ defmodule GraphQLWSClient.Iterator do
 
   @impl true
   def terminate(_reason, %State{} = state) do
-    Process.demonitor(state.monitor_ref)
+    if state.monitor_ref do
+      Process.demonitor(state.monitor_ref, [:flush])
+    end
 
     if state.subscription_id do
       GraphQLWSClient.unsubscribe(state.client, state.subscription_id)

test/graphql_ws_client/iterator_test.exs

@@ -9,6 +9,7 @@ defmodule GraphQLWSClient.IteratorTest do
   alias GraphQLWSClient.Iterator
   alias GraphQLWSClient.Iterator.Opts
   alias GraphQLWSClient.Message
+  alias GraphQLWSClient.SocketError
 
   setup :set_mox_from_context
   setup :verify_on_exit!
@@ -21,19 +22,18 @@ defmodule GraphQLWSClient.IteratorTest do
   @payload_2 %{"baz" => 23}
 
   setup do
-    test_pid = self()
-
-    expect(MockDriver, :connect, fn @conn ->
-      send(test_pid, :connected)
-      {:ok, @conn}
-    end)
-
-    client = start_supervised!({GraphQLWSClient, @config}, id: :test_client)
-    {:ok, client: client, test_pid: test_pid}
+    {:ok, test_pid: self()}
   end
 
   describe "open!/4" do
-    test "success", %{client: client, test_pid: test_pid} do
+    test "success", %{test_pid: test_pid} do
+      expect(MockDriver, :connect, fn @conn ->
+        send(test_pid, :connected)
+        {:ok, @conn}
+      end)
+
+      client = start_supervised!({GraphQLWSClient, @config}, id: :test_client)
+
       expect(MockDriver, :push_message, fn @conn,
                                            %Message{
                                              type: :subscribe,
@@ -51,15 +51,29 @@ defmodule GraphQLWSClient.IteratorTest do
       assert_receive :subscribed
     end
 
-    test "non-numeric buffer size", %{client: client} do
+    test "non-numeric buffer size", %{test_pid: test_pid} do
+      expect(MockDriver, :connect, fn @conn ->
+        send(test_pid, :connected)
+        {:ok, @conn}
+      end)
+
+      client = start_supervised!({GraphQLWSClient, @config}, id: :test_client)
+
       assert_receive :connected
 
       assert_raise ArgumentError, "invalid buffer size", fn ->
         Iterator.open!(client, @query, @variables, buffer_size: "__invalid__")
       end
     end
 
-    test "negative buffer size", %{client: client} do
+    test "negative buffer size", %{test_pid: test_pid} do
+      expect(MockDriver, :connect, fn @conn ->
+        send(test_pid, :connected)
+        {:ok, @conn}
+      end)
+
+      client = start_supervised!({GraphQLWSClient, @config}, id: :test_client)
+
       assert_receive :connected
 
       assert_raise ArgumentError, "invalid buffer size", fn ->
@@ -69,8 +83,16 @@ defmodule GraphQLWSClient.IteratorTest do
   end
 
   describe "next/1" do
-    setup %{client: client} do
+    setup %{test_pid: test_pid} do
+      expect(MockDriver, :connect, fn @conn ->
+        send(test_pid, :connected)
+        {:ok, @conn}
+      end)
+
+      client = start_supervised!({GraphQLWSClient, @config}, id: :test_client)
+
       {:ok,
+       client: client,
        opts: Opts.new(client: client, query: @query, variables: @variables)}
     end