Add abillity to replace device in LiveScreen (#28)

* Add abillity to replace device in LiveScreen

Signed-off-by: Connor Rigby <connorr@hey.com>

* Add tests for set_device and get_device

---------

Signed-off-by: Connor Rigby <connorr@hey.com>
Co-authored-by: Artur Plysyuk <ifuelen@gmail.com>

lib/owl/live_screen.ex

@@ -78,6 +78,29 @@ defmodule Owl.LiveScreen do
     GenServer.start_link(__MODULE__, opts, server_options)
   end
 
+  @doc """
+  Allows updating the device of the LiveScreen.
+
+  This is useful for example when using the Erlang SSH console.
+
+  ## Example
+      group_leader = Process.group_leader()
+
+      Owl.LiveScreen.set_device(group_leader)
+  """
+  @spec set_device(GenServer.server(), IO.device()) :: :ok
+  def set_device(server \\ __MODULE__, new_device) do
+    GenServer.call(server, {:set_device, new_device})
+  end
+
+  @doc """
+  Returns the current device used by the LiveScreen.
+  """
+  @spec get_device(GenServer.server()) :: IO.device()
+  def get_device(server \\ __MODULE__) do
+    GenServer.call(server, :get_device)
+  end
+
   @doc """
   Redirects output from `:stdio` to `#{inspect(__MODULE__)}`.
 
@@ -282,6 +305,18 @@ defmodule Owl.LiveScreen do
   end
 
   @impl true
+  def handle_call({:set_device, new_device}, _, state) do
+    state = render(state)
+
+    state = init_state(state.terminal_width, state.refresh_every, new_device)
+
+    {:reply, :ok, state}
+  end
+
+  def handle_call(:get_device, _, state) do
+    {:reply, state.device, state}
+  end
+
   def handle_call(:flush, _, state) do
     state = render(state)
 

test/owl/live_screen_test.exs

@@ -100,4 +100,30 @@ defmodule Owl.LiveScreenTest do
       terminal_height: 5
     )
   end
+
+  test "set_device and get_device" do
+    terminal_width = 50
+    terminal_height = 20
+
+    device1 =
+      ExUnit.Callbacks.start_supervised!(
+        {VirtualLiveScreen.Device, pid: self(), columns: terminal_width, rows: terminal_height},
+        id: :device1
+      )
+
+    device2 =
+      ExUnit.Callbacks.start_supervised!(
+        {VirtualLiveScreen.Device, pid: self(), columns: terminal_width, rows: terminal_height},
+        id: :device2
+      )
+
+    live_screen_pid =
+      ExUnit.Callbacks.start_supervised!(
+        {Owl.LiveScreen, device: device1, terminal_width: terminal_width}
+      )
+
+    assert Owl.LiveScreen.get_device(live_screen_pid) == device1
+    assert Owl.LiveScreen.set_device(live_screen_pid, device2) == :ok
+    assert Owl.LiveScreen.get_device(live_screen_pid) == device2
+  end
 end