1 files changed, 13 insertions, 22 deletions
diff --git a/lib/polyjuice/client/sync.ex b/lib/polyjuice/client/sync.ex
index 70f357c..77e66b6 100644
--- a/lib/polyjuice/client/sync.ex
+++ b/lib/polyjuice/client/sync.ex
@@ -23,13 +23,13 @@ defmodule Polyjuice.Client.Sync do
   """
   @spec start_link([...]) :: {:ok, pid}
   def start_link([client, listener | opts])
-      when is_pid(listener) and is_list(opts) do
+      when (is_pid(listener) or is_function(listener)) and is_list(opts) do
     Task.start_link(__MODULE__, :sync, [client, listener, opts])
   end
 
-  @enforce_keys [:listener, :access_token, :homeserver_url, :uri, :user_id, :storage]
+  @enforce_keys [:send, :access_token, :homeserver_url, :uri, :user_id, :storage]
   defstruct [
-    :listener,
+    :send,
     :conn_ref,
     :access_token,
     :homeserver_url,
@@ -96,7 +96,7 @@ defmodule Polyjuice.Client.Sync do
     uri = URI.merge(homeserver_url, @sync_path)
 
     connect(%__MODULE__{
-      listener: listener,
+      send: if(is_function(listener), do: listener, else: &send(listener, &1)),
       access_token: access_token,
       homeserver_url: homeserver_url,
       user_id: user_id,
@@ -122,7 +122,7 @@ defmodule Polyjuice.Client.Sync do
          ) do
       {:ok, conn_ref} ->
         Logger.info("Connected to sync")
-        send(state.listener, {:connected})
+        state.send({:connected})
 
         if state.set_filter do
           set_filter(%{state | conn_ref: conn_ref, backoff: nil})
@@ -203,7 +203,7 @@ defmodule Polyjuice.Client.Sync do
         backoff = calc_backoff(state.backoff)
         Logger.error("Set filter error: closed; retrying in #{backoff} seconds.")
         connect(%{state | backoff: backoff, conn_ref: nil})
-        send(state.listener, {:disconnected})
+        state.send({:disconnected})
 
       # FIXME: what other error codes do we need to handle?
       {:error, err} ->
@@ -240,7 +240,7 @@ defmodule Polyjuice.Client.Sync do
               Polyjuice.Client.Storage.set_sync_token(state.storage, next_batch)
 
               if not state.initial_done do
-                send(state.listener, {:initial_sync_completed})
+                state.send({:initial_sync_completed})
               end
 
               do_sync(%{state | since: next_batch, backoff: nil, initial_done: true})
@@ -267,7 +267,7 @@ defmodule Polyjuice.Client.Sync do
         backoff = calc_backoff(state.backoff)
         Logger.error("Sync error: closed; retrying in #{backoff} seconds.")
         connect(%{state | backoff: backoff, conn_ref: nil})
-        send(state.listener, {:disconnected})
+        state.send({:disconnected})
 
       # FIXME: what other error codes do we need to handle?
       {:error, err} ->
@@ -293,7 +293,7 @@ defmodule Polyjuice.Client.Sync do
     |> Map.get("leave", [])
     |> Enum.each(fn {k, v} ->
       process_room(k, v, state)
-      send(state.listener, {:left, k})
+      state.send({:left, k})
     end)
   end
 
@@ -302,7 +302,7 @@ defmodule Polyjuice.Client.Sync do
 
     if Map.get(timeline, "limited", false) do
       with {:ok, prev_batch} <- Map.get(timeline, "prev_batch") do
-        send(state.listener, {:limited, roomname, prev_batch})
+        state.send({:limited, roomname, prev_batch})
       end
     end
 
@@ -323,10 +323,7 @@ defmodule Polyjuice.Client.Sync do
          roomname,
          state
        ) do
-    send(
-      state.listener,
-      {:state, roomname, event}
-    )
+    state.send({:state, roomname, event})
 
     state
   end
@@ -336,10 +333,7 @@ defmodule Polyjuice.Client.Sync do
          roomname,
          state
        ) do
-    send(
-      state.listener,
-      {:message, roomname, event}
-    )
+    state.send({:message, roomname, event})
 
     state
   end
@@ -376,10 +370,7 @@ defmodule Polyjuice.Client.Sync do
       |> Map.get("sender")
 
     if inviter do
-      send(
-        state.listener,
-        {:invite, roomname, inviter, invite_state}
-      )
+      state.send({:invite, roomname, inviter, invite_state})
     end
 
     state