From b4696710e7b2465ac04d70065fbf2f49b36289e2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Valim?= <jose.valim@dashbit.co>
Date: Sat, 24 Dec 2022 11:35:06 -0300
Subject: [PATCH] Encapsulate server access (#1605)

---
 lib/livebook/hubs/enterprise_client.ex                | 11 +++--------
 lib/livebook/web_socket/server.ex                     |  2 +-
 lib/livebook_web/live/hub/new/enterprise_component.ex |  5 +----
 3 files changed, 5 insertions(+), 13 deletions(-)

diff --git a/lib/livebook/hubs/enterprise_client.ex b/lib/livebook/hubs/enterprise_client.ex
index 66aeb3b99..37680dc0c 100644
--- a/lib/livebook/hubs/enterprise_client.ex
+++ b/lib/livebook/hubs/enterprise_client.ex
@@ -20,9 +20,9 @@ defmodule Livebook.Hubs.EnterpriseClient do
   @doc """
   Gets the WebSocket server PID.
   """
-  @spec get_server(pid()) :: pid()
-  def get_server(pid) do
-    GenServer.call(pid, :get_server)
+  @spec send_request(pid(), WebSocket.proto()) :: {atom(), term()}
+  def send_request(pid, %_struct{} = data) do
+    Server.send_request(GenServer.call(pid, :get_server), data)
   end
 
   @doc """
@@ -58,11 +58,6 @@ defmodule Livebook.Hubs.EnterpriseClient do
     Phoenix.PubSub.unsubscribe(Livebook.PubSub, @pubsub_topic)
   end
 
-  @doc """
-
-  """
-  @spec broadcast_message(any()) :: :ok
-
   ## GenServer callbacks
 
   @impl true
diff --git a/lib/livebook/web_socket/server.ex b/lib/livebook/web_socket/server.ex
index d14093b1b..7e7b25717 100644
--- a/lib/livebook/web_socket/server.ex
+++ b/lib/livebook/web_socket/server.ex
@@ -40,7 +40,7 @@ defmodule Livebook.WebSocket.Server do
   @doc """
   Sends a Request to given WebSocket Server.
   """
-  @spec send_request(pid(), WebSocket.proto()) :: {:ok, non_neg_integer()}
+  @spec send_request(pid(), WebSocket.proto()) :: {atom(), term()}
   def send_request(conn, %_struct{} = data) do
     Connection.call(conn, {:request, data}, @timeout)
   end
diff --git a/lib/livebook_web/live/hub/new/enterprise_component.ex b/lib/livebook_web/live/hub/new/enterprise_component.ex
index db2ef8ae1..91708e4ab 100644
--- a/lib/livebook_web/live/hub/new/enterprise_component.ex
+++ b/lib/livebook_web/live/hub/new/enterprise_component.ex
@@ -5,7 +5,6 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponent do
 
   alias Livebook.EctoTypes.HexColor
   alias Livebook.Hubs.{Enterprise, EnterpriseClient}
-  alias Livebook.WebSocket.Server
 
   @impl true
   def update(assigns, socket) do
@@ -127,9 +126,7 @@ defmodule LivebookWeb.Hub.New.EnterpriseComponent do
         session_request =
           LivebookProto.SessionRequest.new!(app_version: Livebook.Config.app_version())
 
-        conn = EnterpriseClient.get_server(pid)
-
-        case Server.send_request(conn, session_request) do
+        case EnterpriseClient.send_request(pid, session_request) do
           {:session, session_response} ->
             base = %{base | external_id: session_response.user.id}
             changeset = Enterprise.change_hub(base)