Migrates Livebook Teams HTTP client to Req (#2492)

lib/livebook/teams/requests.ex

@@ -5,7 +5,6 @@ defmodule Livebook.Teams.Requests do
   alias Livebook.Secrets.Secret
   alias Livebook.Teams
   alias Livebook.Teams.{AgentKey, DeploymentGroup, Org}
-  alias Livebook.Utils.HTTP
 
   @doc """
   Send a request to Livebook Team API to create a new org.
@@ -236,7 +235,42 @@ defmodule Livebook.Teams.Requests do
     value |> Ecto.Changeset.change() |> add_errors(struct.__schema__(:fields), errors_map)
   end
 
-  defp auth_headers(team) do
+  defp post(path, json, team \\ nil) do
+    build_req()
+    |> add_team_auth(team)
+    |> request(method: :post, url: path, json: json)
+    |> dispatch_messages(team)
+  end
+
+  defp put(path, json, team) do
+    build_req()
+    |> add_team_auth(team)
+    |> request(method: :put, url: path, json: json)
+    |> dispatch_messages(team)
+  end
+
+  defp delete(path, json, team) do
+    build_req()
+    |> add_team_auth(team)
+    |> request(method: :delete, url: path, json: json)
+    |> dispatch_messages(team)
+  end
+
+  defp get(path, params \\ %{}) do
+    build_req()
+    |> request(method: :get, url: path, params: params)
+  end
+
+  defp build_req() do
+    Req.new(
+      base_url: Livebook.Config.teams_url(),
+      headers: [{"x-lb-version", Livebook.Config.app_version()}]
+    )
+  end
+
+  defp add_team_auth(req, nil), do: req
+
+  defp add_team_auth(req, team) do
     token =
       if team.user_id do
         "#{team.user_id}:#{team.org_id}:#{team.org_key_id}:#{team.session_token}"
@@ -244,60 +278,25 @@ defmodule Livebook.Teams.Requests do
         "#{team.session_token}:#{Livebook.Config.agent_name()}:#{team.org_id}:#{team.org_key_id}"
       end
 
-    [
-      {"x-lb-version", Livebook.Config.app_version()},
-      {"authorization", "Bearer " <> token}
-    ]
+    Req.Request.merge_options(req, auth: {:bearer, token})
   end
 
-  defp post(path, json, team \\ nil) do
-    body = {"application/json", Jason.encode!(json)}
-    headers = if team, do: auth_headers(team), else: []
-
-    request(:post, path, body: body, headers: headers)
-    |> dispatch_messages(team)
-  end
-
-  defp put(path, json, team) do
-    body = {"application/json", Jason.encode!(json)}
-
-    request(:put, path, body: body, headers: auth_headers(team))
-    |> dispatch_messages(team)
-  end
-
-  defp delete(path, json, team) do
-    body = {"application/json", Jason.encode!(json)}
-
-    request(:delete, path, body: body, headers: auth_headers(team))
-    |> dispatch_messages(team)
-  end
-
-  defp get(path, params \\ %{}) do
-    query_string = URI.encode_query(params)
-    path = if query_string != "", do: "#{path}?#{query_string}", else: path
-
-    request(:get, path, headers: [])
-  end
-
-  defp request(method, path, opts) do
-    endpoint = Livebook.Config.teams_url()
-    url = endpoint <> path
-
-    case HTTP.request(method, url, opts) do
-      {:ok, 204, _headers, body} ->
+  defp request(req, opts) do
+    case Req.request(req, opts) do
+      {:ok, %{status: 204, body: body}} ->
         {:ok, body}
 
-      {:ok, status, headers, body} when status in 200..299 ->
-        if json?(headers),
-          do: {:ok, Jason.decode!(body)},
-          else: {:error, body}
+      {:ok, %{status: status} = response} when status in 200..299 ->
+        if json?(response),
+          do: {:ok, response.body},
+          else: {:error, response.body}
 
-      {:ok, status, headers, body} when status in [410, 422] ->
-        if json?(headers),
-          do: {:error, Jason.decode!(body)},
-          else: {:transport_error, body}
+      {:ok, %{status: status} = response} when status in [410, 422] ->
+        if json?(response),
+          do: {:error, response.body},
+          else: {:transport_error, response.body}
 
-      {:ok, 401, _headers, _body} ->
+      {:ok, %{status: 401}} ->
         {:transport_error,
          "You are not authorized to perform this action, make sure you have the access or you are not in a Livebook Agent instance"}
 
@@ -324,7 +323,7 @@ defmodule Livebook.Teams.Requests do
 
   defp dispatch_messages(result, _), do: result
 
-  defp json?(headers) do
-    HTTP.fetch_content_type(headers) == {:ok, "application/json"}
+  defp json?(response) do
+    "application/json; charset=utf-8" in Req.Response.get_header(response, "content-type")
   end
 end

test/livebook_teams/hubs/team_client_test.exs

@@ -112,7 +112,8 @@ defmodule Livebook.Hubs.TeamClientTest do
       # receives `{:event, :deployment_group_created, deployment_group}` event
       assert_receive {:deployment_group_created, %{name: ^name, mode: ^mode} = deployment_group}
 
-      updated_deployment_group = %{deployment_group | mode: :offline}
+      new_name = "ChonkyCat123"
+      updated_deployment_group = %{deployment_group | name: new_name}
 
       assert {:ok, ^id} =
                Livebook.Teams.update_deployment_group(
@@ -121,10 +122,7 @@ defmodule Livebook.Hubs.TeamClientTest do
                )
 
       # receives `{:deployment_group_updated, deployment_group}` event
-      assert_receive {:deployment_group_updated, ^updated_deployment_group}
-
-      # receives `{:deployment_group_deleted, deployment_group}` event
-      assert_receive {:deployment_group_deleted, ^updated_deployment_group}
+      assert_receive {:deployment_group_updated, %{name: ^new_name, mode: ^mode}}
     end
 
     test "receives the agent key events", %{team: team} do
@@ -135,21 +133,23 @@ defmodule Livebook.Hubs.TeamClientTest do
       id = to_string(id)
 
       # receives `{:event, :deployment_group_created, :deployment_group}` event
-      assert_receive {:deployment_group_created, %{id: ^id} = deployment_group}
+      assert_receive {:deployment_group_created,
+                      %{id: ^id, agent_keys: [built_in_agent_key]} = deployment_group}
 
       # creates the agent key
       assert Livebook.Teams.create_agent_key(team, deployment_group) == :ok
 
       # since the `agent_key` belongs to a deployment group,
       # we dispatch the `{:event, :deployment_group_updated, :deployment_group}` event
-      assert_receive {:deployment_group_updated, %{id: ^id, agent_keys: [agent_key]}}
+      assert_receive {:deployment_group_updated,
+                      %{id: ^id, agent_keys: [^built_in_agent_key, agent_key]}}
 
       # deletes the agent key
       assert Livebook.Teams.delete_agent_key(team, agent_key) == :ok
 
       # since the `agent_key` belongs to a deployment group,
       # we dispatch the `{:event, :deployment_group_updated, :deployment_group}` event
-      assert_receive {:deployment_group_updated, %{id: ^id, agent_keys: []}}
+      assert_receive {:deployment_group_updated, %{id: ^id, agent_keys: [^built_in_agent_key]}}
     end
   end
 

test/livebook_teams/teams_test.exs

@@ -220,26 +220,5 @@ defmodule Livebook.TeamsTest do
 
       assert "can't be blank" in errors_on(changeset).name
     end
-
-    test "returns changeset errors when the new mode is invalid", %{user: user, node: node} do
-      team = create_team_hub(user, node)
-      deployment_group = build(:deployment_group, name: "BAR", mode: :online)
-
-      assert {:ok, id} = Teams.create_deployment_group(team, deployment_group)
-
-      update_deployment_group = %{deployment_group | id: to_string(id), mode: nil}
-
-      assert {:error, changeset} =
-               Teams.update_deployment_group(team, update_deployment_group)
-
-      assert "can't be blank" in errors_on(changeset).mode
-
-      update_deployment_group = %{deployment_group | id: to_string(id), mode: :invalid}
-
-      assert {:error, changeset} =
-               Teams.update_deployment_group(team, update_deployment_group)
-
-      assert "is invalid" in errors_on(changeset).mode
-    end
   end
 end

test/livebook_teams/web/hub/deployment_group_live_test.exs

@@ -197,16 +197,14 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
   end
 
   test "updates an existing secret", %{conn: conn, hub: hub} do
-    insert_deployment_group(
-      name: "TEAMS_EDIT_DEPLOYMENT_GROUP",
-      mode: :online,
-      hub_id: hub.id
-    )
+    name = "TEAMS_EDIT_DEPLOYMENT_GROUP"
+    mode = :online
+    insert_deployment_group(name: name, mode: mode, hub_id: hub.id)
 
     hub_id = hub.id
 
     assert_receive {:deployment_group_created,
-                    %DeploymentGroup{name: "TEAMS_EDIT_DEPLOYMENT_GROUP", hub_id: ^hub_id} =
+                    %DeploymentGroup{name: ^name, mode: ^mode, hub_id: ^hub_id, secrets: []} =
                       deployment_group}
 
     id = deployment_group.id
@@ -317,16 +315,19 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
   end
 
   test "creates an agent key", %{conn: conn, hub: hub} do
-    insert_deployment_group(
-      name: "TEAMS_AGENT_KEY_DEPLOYMENT_GROUP",
-      mode: :online,
-      hub_id: hub.id
-    )
+    name = "TEAMS_AGENT_KEY_DEPLOYMENT_GROUP"
+    mode = :online
+    insert_deployment_group(name: name, mode: mode, hub_id: hub.id)
 
     hub_id = hub.id
 
     assert_receive {:deployment_group_created,
-                    %DeploymentGroup{name: "TEAMS_AGENT_KEY_DEPLOYMENT_GROUP", hub_id: ^hub_id} =
+                    %DeploymentGroup{
+                      name: ^name,
+                      mode: ^mode,
+                      agent_keys: [built_in_agent_key],
+                      hub_id: ^hub_id
+                    } =
                       deployment_group}
 
     id = deployment_group.id
@@ -341,24 +342,26 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
     render_confirm(view)
 
     assert_receive {:deployment_group_updated,
-                    %Livebook.Teams.DeploymentGroup{id: ^id, agent_keys: [agent_key]}}
+                    %Livebook.Teams.DeploymentGroup{
+                      id: ^id,
+                      agent_keys: [^built_in_agent_key, agent_key]
+                    }}
 
     assert render(view) =~ agent_key.key
   end
 
   test "deletes an agent key", %{conn: conn, hub: hub} do
-    insert_agent_key(
-      name: "TEAMS_AGENT_KEY_DEPLOYMENT_GROUP_DELETE",
-      mode: :online,
-      hub_id: hub.id
-    )
+    name = "TEAMS_AGENT_KEY_DEPLOYMENT_GROUP_DELETE"
+    mode = :online
+    insert_deployment_group(name: name, mode: mode, hub_id: hub.id)
 
     hub_id = hub.id
 
-    assert_receive {:deployment_group_updated,
+    assert_receive {:deployment_group_created,
                     %DeploymentGroup{
                       id: id,
-                      name: "TEAMS_AGENT_KEY_DEPLOYMENT_GROUP_DELETE",
+                      name: ^name,
+                      mode: ^mode,
                       hub_id: ^hub_id,
                       agent_keys: [%{deployment_group_id: id} = agent_key]
                     }}
@@ -381,20 +384,14 @@ defmodule LivebookWeb.Integration.Hub.DeploymentGroupLiveTest do
 
   test "doesn't show agent key section for offline deployment groups",
        %{conn: conn, hub: hub} do
-    insert_deployment_group(
-      name: "TEAMS_AGENT_KEY_DEPLOYMENT_GROUP",
-      mode: :online,
-      hub_id: hub.id
-    )
+    name = "TEAMS_AGENT_KEY_DEPLOYMENT_GROUP_DELETE"
+    mode = :offline
+    insert_deployment_group(name: name, mode: mode, hub_id: hub.id)
 
     hub_id = hub.id
 
     assert_receive {:deployment_group_created,
-                    %DeploymentGroup{
-                      id: id,
-                      name: "TEAMS_AGENT_KEY_DEPLOYMENT_GROUP",
-                      hub_id: ^hub_id
-                    }}
+                    %DeploymentGroup{id: id, name: ^name, mode: ^mode, hub_id: ^hub_id}}
 
     {:ok, view, _html} = live(conn, ~p"/hub/#{hub.id}/deployment-groups/edit/#{id}")
 

test/livebook_teams/web/hub/edit_live_test.exs

@@ -355,14 +355,13 @@ defmodule LivebookWeb.Integration.Hub.EditLiveTest do
     end
 
     test "updates an existing deployment group", %{conn: conn, hub: hub} do
-      insert_deployment_group(
-        name: "TEAM_EDIT_DEPLOYMENT_GROUP",
-        mode: :online,
-        hub_id: hub.id
-      )
+      name = "TEAM_EDIT_DEPLOYMENT_GROUP"
+      mode = :online
+      insert_deployment_group(name: name, mode: mode, hub_id: hub.id)
 
       assert_receive {:deployment_group_created,
-                      %DeploymentGroup{name: "TEAM_EDIT_DEPLOYMENT_GROUP"} = deployment_group}
+                      %DeploymentGroup{name: ^name, mode: ^mode, agent_keys: [_]} =
+                        deployment_group}
 
       {:ok, view, _html} = live(conn, ~p"/hub/#{hub.id}")
 

test/support/factory.ex

@@ -114,16 +114,6 @@ defmodule Livebook.Factory do
     %{deployment_group | id: to_string(id)}
   end
 
-  def insert_agent_key(attrs \\ %{}) do
-    deployment_group = build(:deployment_group, attrs)
-    hub = Livebook.Hubs.fetch_hub!(deployment_group.hub_id)
-    {:ok, id} = Livebook.Teams.create_deployment_group(hub, deployment_group)
-    deployment_group = %{deployment_group | id: to_string(id)}
-    :ok = Livebook.Teams.create_agent_key(hub, deployment_group)
-
-    deployment_group
-  end
-
   def insert_env_var(factory_name, attrs \\ %{}) do
     env_var = build(factory_name, attrs)
     attributes = env_var |> Map.from_struct() |> Map.to_list()