handle updating reg state for structs and maps

diff --git a/lib/client/handler.ex b/lib/client/handler.ex
index e4def16c1d952643e40cf2f3303cfed8360c157c..86237ce9870f21b4f5f47c9f462d35f3bad5b43b 100644 (file)
--- a/lib/client/handler.ex
+++ b/lib/client/handler.ex
@@ -34,7 +34,7 @@ defmodule Wampex.Client.Handler do
 
       @impl true
       def handle_continue({:registered, registrations}, state),
-        do: {:noreply, %{state | registrations: registrations}}
+        do: {:noreply, add_regs(state, registrations)}
 
       @impl true
       def terminate(_reason, %{
@@ -50,7 +50,7 @@ defmodule Wampex.Client.Handler do
         end)
       end
 
-      def do_init(_opts), do: %{client_name: nil}
+      def do_init(_opts), do: %{client_name: nil, registrations: nil}
 
       def handle_invocation_block({:ok, al, kw, state}, id) do
         Client.yield(state.client_name, %Yield{request_id: id, arg_list: al, arg_kw: kw})
@@ -68,6 +68,12 @@ defmodule Wampex.Client.Handler do
         state
       end
 
+      defp add_regs(state, regs) when is_struct(state) do
+        struct(state, %{registrations: regs})
+      end
+
+      defp add_regs(state, regs), do: Map.put(state, :registrations, regs)
+
       defp add_client(state, cn) when is_struct(state) do
         struct(state, %{client_name: cn})
       end