From: Christopher Coté <chris@entropealabs.com>
Date: Tue, 23 Nov 2021 17:52:04 +0000 (-0500)
Subject: handle updating reg state for structs and maps
X-Git-Url: http://git.entropealabs.com/?a=commitdiff_plain;h=6c00747ffd49cbadfaf851f0c2b6c46816fac86e;p=wampex_client.git

handle updating reg state for structs and maps
---

diff --git a/lib/client/handler.ex b/lib/client/handler.ex
index e4def16..86237ce 100644
--- 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