From 88e34076b59ff142f3429ba60641fc72904e9aaa Mon Sep 17 00:00:00 2001
From: Tobias <tobias.casper@gmail.com>
Date: Fri, 29 Mar 2024 17:26:13 +0100
Subject: [PATCH] fix: trap exits in iterator server

---
 lib/graphql_ws_client/iterator.ex        | 10 +++--
 test/graphql_ws_client/iterator_test.exs | 48 +++++++++++++++++-------
 2 files changed, 41 insertions(+), 17 deletions(-)

diff --git a/lib/graphql_ws_client/iterator.ex b/lib/graphql_ws_client/iterator.ex
index 9dc59b7..4701b8c 100644
--- a/lib/graphql_ws_client/iterator.ex
+++ b/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)
diff --git a/test/graphql_ws_client/iterator_test.exs b/test/graphql_ws_client/iterator_test.exs
index 7730a46..349def8 100644
--- a/test/graphql_ws_client/iterator_test.exs
+++ b/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,7 +51,14 @@ 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 ->
@@ -59,7 +66,14 @@ defmodule GraphQLWSClient.IteratorTest do
       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