Deploy notebooks from CLI (#3034)

lib/livebook.ex

@@ -157,6 +157,10 @@ defmodule Livebook do
       config :livebook, teams_url: url, warn_on_live_teams_server: false
     end
 
+    if System.get_env("LIVEBOOK_TEAMS_AUTH") do
+      config :livebook, :persist_storage, false
+    end
+
     if Livebook.Config.boolean!("LIVEBOOK_SHUTDOWN_ENABLED", false) do
       config :livebook, :shutdown_callback, {System, :stop, []}
     end

lib/livebook/hubs.ex

@@ -311,6 +311,7 @@ defmodule Livebook.Hubs do
   @spec get_app_specs() :: list(Livebook.Apps.AppSpec.t())
   def get_app_specs() do
     for hub <- get_hubs(),
+        Provider.connection_spec(hub),
         app_spec <- Provider.get_app_specs(hub),
         do: app_spec
   end

lib/livebook/hubs/team.ex

@@ -94,22 +94,17 @@ defmodule Livebook.Hubs.Team do
       changeset
     end
   end
-
-  @doc """
-  Returns the public key prefix
-  """
-  @spec public_key_prefix() :: String.t()
-  def public_key_prefix(), do: "lb_opk_"
 end
 
 defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   alias Livebook.Hubs.{Team, TeamClient}
-  alias Livebook.Teams.Requests
+  alias Livebook.Teams
   alias Livebook.FileSystem
   alias Livebook.Secrets.Secret
 
-  @teams_key_prefix Livebook.Teams.Org.teams_key_prefix()
+  @teams_key_prefix Teams.Constants.teams_key_prefix()
-  @public_key_prefix Team.public_key_prefix()
+  @public_key_prefix Teams.Constants.public_key_prefix()
+  @deploy_key_prefix Teams.Constants.deploy_key_prefix()
 
   def load(team, fields) do
     {offline?, fields} = Map.pop(fields, :offline?, false)
@@ -137,6 +132,7 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
 
   def type(_team), do: "team"
 
+  def connection_spec(%{session_token: @deploy_key_prefix <> _}), do: nil
   def connection_spec(team), do: {TeamClient, team}
 
   def disconnect(team), do: TeamClient.stop(team.id)
@@ -160,7 +156,7 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
     @teams_key_prefix <> teams_key = team.teams_key
     token = Livebook.Stamping.chapoly_encrypt(metadata, notebook_source, teams_key)
 
-    case Requests.org_sign(team, token) do
+    case Teams.Requests.org_sign(team, token) do
       {:ok, %{"signature" => token_signature}} ->
         stamp = %{"version" => 1, "token" => token, "token_signature" => token_signature}
         {:ok, stamp}
@@ -201,7 +197,7 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   def get_secrets(team), do: TeamClient.get_secrets(team.id)
 
   def create_secret(%Team{} = team, %Secret{} = secret) do
-    case Requests.create_secret(team, secret) do
+    case Teams.Requests.create_secret(team, secret) do
       {:ok, %{"id" => _}} -> :ok
       {:error, %{"errors" => errors}} -> {:error, parse_secret_errors(errors)}
       any -> any
@@ -209,7 +205,7 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   end
 
   def update_secret(%Team{} = team, %Secret{} = secret) do
-    case Requests.update_secret(team, secret) do
+    case Teams.Requests.update_secret(team, secret) do
       {:ok, %{"id" => _}} -> :ok
       {:error, %{"errors" => errors}} -> {:error, parse_secret_errors(errors)}
       any -> any
@@ -217,7 +213,7 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   end
 
   def delete_secret(%Team{} = team, %Secret{} = secret) do
-    case Requests.delete_secret(team, secret) do
+    case Teams.Requests.delete_secret(team, secret) do
       {:ok, _} -> :ok
       {:error, %{"errors" => errors}} -> {:error, parse_secret_errors(errors)}
       any -> any
@@ -227,7 +223,7 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   def get_file_systems(team), do: TeamClient.get_file_systems(team.id)
 
   def create_file_system(%Team{} = team, file_system) do
-    case Requests.create_file_system(team, file_system) do
+    case Teams.Requests.create_file_system(team, file_system) do
       {:ok, %{"id" => _}} -> :ok
       {:error, %{"errors" => errors}} -> {:error, parse_file_system_errors(file_system, errors)}
       any -> any
@@ -235,7 +231,7 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   end
 
   def update_file_system(%Team{} = team, file_system) do
-    case Requests.update_file_system(team, file_system) do
+    case Teams.Requests.update_file_system(team, file_system) do
       {:ok, %{"id" => _}} -> :ok
       {:error, %{"errors" => errors}} -> {:error, parse_file_system_errors(file_system, errors)}
       any -> any
@@ -243,7 +239,7 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   end
 
   def delete_file_system(%Team{} = team, file_system) do
-    case Requests.delete_file_system(team, file_system) do
+    case Teams.Requests.delete_file_system(team, file_system) do
       {:ok, _} -> :ok
       {:error, %{"errors" => errors}} -> {:error, parse_file_system_errors(file_system, errors)}
       any -> any
@@ -264,12 +260,12 @@ defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
   end
 
   defp parse_secret_errors(errors_map) do
-    Requests.to_error_list(Secret, errors_map)
+    Teams.Requests.to_error_list(Secret, errors_map)
   end
 
   defp parse_file_system_errors(%struct{} = file_system, errors_map) do
     %{error_field: field} = FileSystem.external_metadata(file_system)
     errors_map = Map.new(errors_map, fn {_key, values} -> {field, values} end)
-    Requests.to_error_list(struct, errors_map)
+    Teams.Requests.to_error_list(struct, errors_map)
   end
 end

lib/livebook/storage.ex

@@ -186,10 +186,11 @@ defmodule Livebook.Storage do
     # in case it is persisting to disk. terminate/2 is still a no-op.
     Process.flag(:trap_exit, true)
 
-    table = load_or_create_table()
+    persist_storage? = Application.get_env(:livebook, :persist_storage, true)
+    table = load_or_create_table(persist_storage?)
     :persistent_term.put(__MODULE__, table)
 
-    {:ok, %{table: table}}
+    {:ok, %{table: table, persist?: persist_storage?}}
   end
 
   @impl true
@@ -220,6 +221,10 @@ defmodule Livebook.Storage do
   end
 
   @impl true
+  def handle_continue(:save_to_file, %{persist?: false} = state) do
+    {:noreply, state}
+  end
+
   def handle_continue(:save_to_file, %{table: table} = state) do
     file_path = String.to_charlist(config_file_path())
     :ok = :ets.tab2file(table, file_path)
@@ -228,7 +233,11 @@ defmodule Livebook.Storage do
 
   defp table_name(), do: :persistent_term.get(__MODULE__)
 
-  defp load_or_create_table() do
+  defp load_or_create_table(false) do
+    :ets.new(__MODULE__, [:protected, :duplicate_bag])
+  end
+
+  defp load_or_create_table(true) do
     tab =
       if path = config_file_path_for_restore() do
         path

lib/livebook/teams.ex

@@ -11,7 +11,7 @@ defmodule Livebook.Teams do
   import Ecto.Changeset,
     only: [add_error: 3, apply_action: 2, apply_action!: 2, get_field: 2]
 
-  @prefix Org.teams_key_prefix()
+  @teams_key_prefix Teams.Constants.teams_key_prefix()
 
   @doc """
   Creates an Org.
@@ -148,7 +148,7 @@ defmodule Livebook.Teams do
   Derives the secret and sign secret from given `teams_key`.
   """
   @spec derive_key(String.t()) :: bitstring()
-  def derive_key(@prefix <> teams_key) do
+  def derive_key(@teams_key_prefix <> teams_key) do
     binary_key = Base.url_decode64!(teams_key, padding: false)
     Plug.Crypto.KeyGenerator.generate(binary_key, "notebook secret", cache: Plug.Crypto.Keys)
   end
@@ -235,6 +235,45 @@ defmodule Livebook.Teams do
     TeamClient.get_environment_variables(team.id)
   end
 
+  @doc """
+  Fetches the CLI session using a deploy key.
+  """
+  @spec fetch_cli_session(map()) ::
+          {:ok, Team.t()} | {:error, String.t()} | {:transport_error, String.t()}
+  def fetch_cli_session(%{session_token: _, teams_key: _} = config) do
+    with {:ok, %{"name" => name} = attrs} <- Requests.fetch_cli_session(config) do
+      id = "team-#{name}"
+
+      hub =
+        Hubs.save_hub(%Team{
+          id: id,
+          hub_name: name,
+          hub_emoji: "🚀",
+          user_id: nil,
+          org_id: attrs["org_id"],
+          org_key_id: attrs["org_key_id"],
+          session_token: config.session_token,
+          teams_key: config.teams_key,
+          org_public_key: attrs["public_key"]
+        })
+
+      {:ok, hub}
+    end
+  end
+
+  @doc """
+  Deploys the given app deployment to given deployment group using a deploy key.
+  """
+  @spec deploy_app_from_cli(Team.t(), Teams.AppDeployment.t(), String.t()) ::
+          {:ok, String.t()} | {:error, map()} | {:transport_error, String.t()}
+  def deploy_app_from_cli(%Team{} = team, %Teams.AppDeployment{} = app_deployment, name) do
+    case Requests.deploy_app_from_cli(team, app_deployment, name) do
+      {:ok, %{"url" => url}} -> {:ok, url}
+      {:error, %{"errors" => errors}} -> {:error, errors}
+      any -> any
+    end
+  end
+
   defp map_teams_field_to_livebook_field(map, teams_field, livebook_field) do
     if value = map[teams_field] do
       Map.put_new(map, livebook_field, value)

lib/livebook/teams/constants.ex

@@ -0,0 +1,25 @@
+defmodule Livebook.Teams.Constants do
+  @doc """
+  Returns the public key prefix
+  """
+  @spec public_key_prefix() :: String.t()
+  def public_key_prefix(), do: "lb_opk_"
+
+  @doc """
+  Returns the Agent Key prefix
+  """
+  @spec agent_key_prefix() :: String.t()
+  def agent_key_prefix, do: "lb_ak_"
+
+  @doc """
+  Returns the Deploy Key prefix
+  """
+  @spec deploy_key_prefix() :: String.t()
+  def deploy_key_prefix, do: "lb_dk_"
+
+  @doc """
+  Returns the Teams Key prefix
+  """
+  @spec teams_key_prefix() :: String.t()
+  def teams_key_prefix(), do: "lb_tk_"
+end

lib/livebook/teams/org.ex

@@ -2,8 +2,6 @@ defmodule Livebook.Teams.Org do
   use Ecto.Schema
   import Ecto.Changeset
 
-  @prefix "lb_tk_"
-
   @type t :: %__MODULE__{
           id: pos_integer() | nil,
           emoji: String.t() | nil,
@@ -31,7 +29,7 @@ defmodule Livebook.Teams.Org do
   @spec teams_key() :: String.t()
   def teams_key() do
     key = :crypto.strong_rand_bytes(@secret_key_size)
-    @prefix <> Base.url_encode64(key, padding: false)
+    Livebook.Teams.Constants.teams_key_prefix() <> Base.url_encode64(key, padding: false)
   end
 
   @doc """
@@ -50,10 +48,4 @@ defmodule Livebook.Teams.Org do
       message: "should only contain lowercase alphanumeric characters and dashes"
     )
   end
-
-  @doc """
-  Returns the teams key prefix
-  """
-  @spec teams_key_prefix() :: String.t()
-  def teams_key_prefix(), do: @prefix
 end

lib/livebook/teams/requests.ex

@@ -5,6 +5,7 @@ defmodule Livebook.Teams.Requests do
   alias Livebook.Secrets.Secret
   alias Livebook.Teams
 
+  @deploy_key_prefix Teams.Constants.deploy_key_prefix()
   @error_message "Something went wrong, try again later or please file a bug if it persists"
   @unauthorized_error_message "You are not authorized to perform this action, make sure you have the access and you are not in a Livebook App Server/Offline instance"
 
@@ -227,6 +228,34 @@ defmodule Livebook.Teams.Requests do
     get("/api/v1/org/identity", %{access_token: access_token}, team)
   end
 
+  @doc """
+  Send a request to Livebook Team API to return a session using a deploy key.
+  """
+  @spec fetch_cli_session(map()) :: api_result()
+  def fetch_cli_session(config) do
+    post("/api/v1/cli/auth", %{}, config)
+  end
+
+  @doc """
+  Send a request to Livebook Team API to deploy an app using a deploy key.
+  """
+  @spec deploy_app_from_cli(Team.t(), Teams.AppDeployment.t(), String.t()) :: api_result()
+  def deploy_app_from_cli(team, app_deployment, deployment_group_name) do
+    secret_key = Teams.derive_key(team.teams_key)
+
+    params = %{
+      title: app_deployment.title,
+      slug: app_deployment.slug,
+      multi_session: app_deployment.multi_session,
+      access_type: app_deployment.access_type,
+      deployment_group_name: deployment_group_name,
+      sha: app_deployment.sha
+    }
+
+    encrypted_content = Teams.encrypt(app_deployment.file, secret_key)
+    upload("/api/v1/cli/org/apps", encrypted_content, params, team)
+  end
+
   @doc """
   Normalizes errors map into errors for the given schema.
   """
@@ -271,7 +300,6 @@ defmodule Livebook.Teams.Requests do
   defp upload(path, content, params, team) do
     build_req(team)
     |> Req.Request.put_header("content-length", "#{byte_size(content)}")
-    |> Req.Request.put_private(:deploy, true)
     |> Req.post(url: path, params: params, body: content)
     |> handle_response()
     |> dispatch_messages(team)
@@ -280,7 +308,6 @@ defmodule Livebook.Teams.Requests do
   defp build_req(team) do
     Req.new(base_url: Livebook.Config.teams_url())
     |> Req.Request.put_new_header("x-lb-version", Livebook.Config.app_version())
-    |> Req.Request.append_response_steps(transform_teams_response: &transform_response/1)
     |> Livebook.Utils.req_attach_defaults()
     |> add_team_auth(team)
   end
@@ -291,6 +318,11 @@ defmodule Livebook.Teams.Requests do
     Req.Request.append_request_steps(req, unauthorized: &{&1, Req.Response.new(status: 401)})
   end
 
+  defp add_team_auth(req, %{session_token: @deploy_key_prefix <> _} = team) do
+    token = "#{team.session_token}:#{Teams.Org.key_hash(%Teams.Org{teams_key: team.teams_key})}"
+    Req.Request.merge_options(req, auth: {:bearer, token})
+  end
+
   defp add_team_auth(req, %{user_id: nil} = team) do
     agent_name = Livebook.Config.agent_name()
     token = "#{team.session_token}:#{agent_name}:#{team.org_id}:#{team.org_key_id}"
@@ -303,17 +335,6 @@ defmodule Livebook.Teams.Requests do
     Req.Request.merge_options(req, auth: {:bearer, token})
   end
 
-  defp transform_response({request, response}) do
-    case {request, response} do
-      {request, %{status: 400, body: %{"errors" => %{"detail" => error}}}}
-      when request.private.deploy ->
-        {request, %{response | status: 422, body: %{"errors" => %{"file" => [error]}}}}
-
-      _otherwise ->
-        {request, response}
-    end
-  end
-
   defp handle_response(response) do
     case response do
       {:ok, %{status: status} = response} when status in 200..299 -> {:ok, response.body}

lib/livebook_cli.ex

@@ -1,19 +1,23 @@
 defmodule LivebookCLI do
-  def usage() do
+  alias LivebookCLI.{Task, Utils}
-    """
+
+  @help_args ["--help", "-h"]
+  @version_args ["--version", "-v"]
+
+  def usage,
+    do: """
     Usage: livebook [command] [options]
 
     Available commands:
 
       livebook server    Starts the Livebook web application
+      livebook deploy    Deploys a notebook to Livebook Teams
 
-    The --help and --version options can be given instead of a command for usage and versioning information.
+    The --help and --version options can be given instead of a command for usage and versioning information.\
     """
-  end
 
   def main(args) do
     {:ok, _} = Application.ensure_all_started(:elixir)
-
     extract_priv!()
 
     :ok = Application.load(:livebook)
@@ -22,62 +26,32 @@ defmodule LivebookCLI do
       Application.put_env(:elixir, :ansi_enabled, true)
     end
 
-    call(args)
+    case args do
+      [arg] when arg in @help_args -> display_help()
+      [arg] when arg in @version_args -> display_version()
+      [name, arg] when arg in @help_args -> Task.usage(name)
+      [name | args] -> Task.call(name, List.delete(args, name))
+      _args -> Utils.log_info(usage())
+    end
   end
 
   defp unix?(), do: match?({:unix, _}, :os.type())
 
-  defp call([arg]) when arg in ["--help", "-h"], do: display_help()
-  defp call([arg]) when arg in ["--version", "-v"], do: display_version()
-
-  defp call([task_name | args]) do
-    case find_task(task_name) do
-      nil ->
-        IO.ANSI.format([:red, "Unknown command #{task_name}\n"]) |> IO.puts()
-        IO.write(usage())
-
-      task ->
-        call_task(task, args)
-    end
-  end
-
-  defp call(_args), do: IO.write(usage())
-
-  defp find_task("server"), do: LivebookCLI.Server
-  defp find_task(_), do: nil
-
-  defp call_task(task, [arg]) when arg in ["--help", "-h"] do
-    IO.write(task.usage())
-  end
-
-  defp call_task(task, args) do
-    try do
-      task.call(args)
-    rescue
-      error in OptionParser.ParseError ->
-        IO.ANSI.format([
-          :red,
-          Exception.message(error),
-          "\n\nFor more information try --help"
-        ])
-        |> IO.puts()
-
-      error ->
-        IO.ANSI.format([:red, Exception.format(:error, error, __STACKTRACE__), "\n"]) |> IO.puts()
-    end
-  end
-
   defp display_help() do
-    IO.puts("Livebook is an interactive notebook system for Elixir\n")
+    Utils.log_info("""
-    IO.write(usage())
+    Livebook is an interactive notebook system for Elixir
+
+    #{usage()}\
+    """)
   end
 
   defp display_version() do
-    IO.puts(:erlang.system_info(:system_version))
+    Utils.log_info("""
-    IO.puts("Elixir " <> System.build_info()[:build])
+    #{:erlang.system_info(:system_version)}
+    Elixir #{System.build_info()[:build]}
 
-    version = Livebook.Config.app_version()
+    Livebook #{Livebook.Config.app_version()}\
-    IO.puts("\nLivebook #{version}")
+    """)
   end
 
   import Record
@@ -104,14 +78,10 @@ defmodule LivebookCLI do
         List.starts_with?(name, in_archive_priv_path)
       end
 
-      case :zip.extract(archive, cwd: String.to_charlist(archive_dir), file_filter: file_filter) do
+      opts = [cwd: String.to_charlist(archive_dir), file_filter: file_filter]
-        {:ok, _} ->
-          :ok
 
-        {:error, error} ->
+      with {:error, error} <- :zip.extract(archive, opts) do
-          print_error_and_exit(
+        raise "Livebook failed to extract archive files, reason: #{inspect(error)}"
-            "Livebook failed to extract archive files, reason: #{inspect(error)}"
-          )
       end
 
       File.touch!(extracted_path)
@@ -120,10 +90,4 @@ defmodule LivebookCLI do
     priv_dir = Path.join(archive_dir, in_archive_priv_path)
     Application.put_env(:livebook, :priv_dir, priv_dir, persistent: true)
   end
-
-  @spec print_error_and_exit(String.t()) :: no_return()
-  defp print_error_and_exit(message) do
-    IO.ANSI.format([:red, message]) |> IO.puts()
-    System.halt(1)
-  end
 end

lib/livebook_cli/deploy.ex

@@ -0,0 +1,208 @@
+defmodule LivebookCLI.Deploy do
+  import LivebookCLI.Utils
+  alias Livebook.Teams
+
+  @behaviour LivebookCLI.Task
+
+  @deploy_key_prefix Teams.Constants.deploy_key_prefix()
+  @teams_key_prefix Teams.Constants.teams_key_prefix()
+
+  @impl true
+  def usage() do
+    """
+    Usage: livebook deploy [options] filename|directory
+
+    ## Available options
+
+      --deploy-key        Sets the deploy key to authenticate with Livebook Teams
+      --teams-key         Sets the Teams key to authenticate with Livebook Teams and encrypt the Livebook app
+      --deployment-group  The deployment group name which you want to deploy to
+
+    The --help option can be given to print this notice.
+
+    ## Examples
+
+    Deploys a single notebook:
+
+        livebook deploy --deploy-key="lb_dk_..." --teams-key="lb_tk_..." --deployment-group "online" path/to/app1.livemd
+
+    Deploys multiple notebooks:
+
+        livebook deploy --deploy-key="lb_dk_..." --teams-key="lb_tk_..." --deployment-group "online" path/to/*.livemd\
+    """
+  end
+
+  @switches [
+    deploy_key: :string,
+    teams_key: :string,
+    deployment_group: :string
+  ]
+
+  @impl true
+  def call(args) do
+    Application.put_env(:livebook, :persist_storage, false)
+    {:ok, _} = Application.ensure_all_started(:livebook)
+    config = config_from_args(args)
+    ensure_config!(config)
+
+    team = authenticate_cli!(config)
+    deploy_to_teams(team, config)
+  end
+
+  defp config_from_args(args) do
+    {opts, paths} = OptionParser.parse!(args, strict: @switches)
+
+    %{
+      paths: paths,
+      session_token: opts[:deploy_key],
+      teams_key: opts[:teams_key],
+      deployment_group: opts[:deployment_group]
+    }
+  end
+
+  defp ensure_config!(config) do
+    log_debug("Validating config from options...")
+
+    errors =
+      Enum.reduce(config, %{}, fn
+        {key, value}, acc when value in ["", nil] ->
+          add_error(acc, normalize_key(key), "can't be blank")
+
+        {:session_token, value}, acc ->
+          if not String.starts_with?(value, @deploy_key_prefix) do
+            add_error(acc, normalize_key(:session_token), "must be a Livebook Teams Deploy Key")
+          else
+            acc
+          end
+
+        {:teams_key, value}, acc ->
+          if not String.starts_with?(value, @teams_key_prefix) do
+            add_error(acc, normalize_key(:teams_key), "must be a Livebook Teams Key")
+          else
+            acc
+          end
+
+        {:paths, values}, acc ->
+          Enum.reduce_while(values, acc, &validate_path/2)
+
+        _otherwise, acc ->
+          acc
+      end)
+
+    if Map.keys(errors) == [] do
+      :ok
+    else
+      raise LivebookCLI.Error, """
+      You configuration is invalid, make sure you are using the correct options for this task.
+
+      #{format_errors(errors, " * ")}\
+      """
+    end
+  end
+
+  defp validate_path(value, acc) do
+    cond do
+      not File.exists?(value) ->
+        {:halt, add_error(acc, normalize_key(:paths), "must be a valid path")}
+
+      File.dir?(value) ->
+        {:halt, add_error(acc, normalize_key(:paths), "must be a file path")}
+
+      true ->
+        {:cont, acc}
+    end
+  end
+
+  defp authenticate_cli!(config) do
+    log_debug("Authenticating CLI...")
+
+    case Teams.fetch_cli_session(config) do
+      {:ok, team} -> team
+      {:error, error} -> raise LivebookCLI.Error, error
+      {:transport_error, error} -> raise LivebookCLI.Error, error
+    end
+  end
+
+  defp deploy_to_teams(team, config) do
+    if length(config.paths) == 1 do
+      log_debug("Found 1 notebook")
+    else
+      log_debug("Found #{length(config.paths)} notebooks")
+    end
+
+    log_info("Deploying notebooks:")
+
+    for path <- config.paths do
+      log_info(" * Preparing to deploy notebook #{Path.basename(path)}")
+      files_dir = Livebook.FileSystem.File.local(path)
+
+      with {:ok, content} <- File.read(path),
+           {:ok, app_deployment} <- prepare_app_deployment(path, content, files_dir) do
+        case Livebook.Teams.deploy_app_from_cli(team, app_deployment, config.deployment_group) do
+          {:ok, url} ->
+            log_info([:green, "  * #{app_deployment.title} deployed successfully. (#{url})"])
+
+          {:error, errors} ->
+            log_error("  * #{app_deployment.title} failed to deploy.")
+            errors = normalize_errors(errors)
+
+            raise LivebookCLI.Error, """
+            #{format_errors(errors, "  * ")}
+
+            #{Teams.Requests.error_message()}\
+            """
+
+          {:transport_error, reason} ->
+            log_error("  * #{app_deployment.title} failed to deploy.")
+            raise LivebookCLI.Error, reason
+        end
+      end
+    end
+
+    :ok
+  end
+
+  defp prepare_app_deployment(path, content, files_dir) do
+    case Livebook.Teams.AppDeployment.new(content, files_dir) do
+      {:ok, app_deployment} ->
+        {:ok, app_deployment}
+
+      {:warning, warnings} ->
+        raise LivebookCLI.Error, """
+        Deployment for notebook #{Path.basename(path)} failed because the notebook has some warnings:
+        #{format_list(warnings, " * ")}
+        """
+
+      {:error, reason} ->
+        raise LivebookCLI.Error, "Failed to handle I/O operations: #{reason}"
+    end
+  end
+
+  defp add_error(errors, key, message) do
+    Map.update(errors, key, [message], &[message | &1])
+  end
+
+  def normalize_errors(%{} = errors) do
+    for {key, values} <- errors, into: %{} do
+      {normalize_key(key), values}
+    end
+  end
+
+  defp normalize_key(key) when is_atom(key), do: to_string(key) |> normalize_key()
+  defp normalize_key("session_token"), do: "Deploy Key"
+  defp normalize_key("teams_key"), do: "Teams Key"
+  defp normalize_key("deployment_group"), do: "Deployment Group"
+  defp normalize_key("paths"), do: "File Paths"
+
+  defp format_errors(errors, prefix) do
+    errors
+    |> Enum.map(fn {key, values} ->
+      values |> Enum.map(&"#{prefix}#{key} #{&1}") |> Enum.join("\n")
+    end)
+    |> Enum.join("\n")
+  end
+
+  defp format_list(errors, prefix) do
+    errors |> Enum.map(&"#{prefix}#{&1}") |> Enum.join("\n")
+  end
+end

lib/livebook_cli/error.ex

@@ -0,0 +1,3 @@
+defmodule LivebookCLI.Error do
+  defexception [:message]
+end

lib/livebook_cli/server.ex

@@ -1,4 +1,5 @@
 defmodule LivebookCLI.Server do
+  import LivebookCLI.Utils
   @behaviour LivebookCLI.Task
 
   @external_resource "README.md"
@@ -63,8 +64,7 @@ defmodule LivebookCLI.Server do
 
     Starts a server and imports the notebook at the given URL:
 
-        livebook server https://example.com/my-notebook.livemd
+        livebook server https://example.com/my-notebook.livemd\
-
     """
   end
 
@@ -84,14 +84,13 @@ defmodule LivebookCLI.Server do
 
         case check_endpoint_availability(base_url) do
           :livebook_running ->
-            IO.puts("Livebook already running on #{base_url}")
+            log_info("Livebook already running on #{base_url}")
             open_from_args(base_url, extra_args)
 
           :taken ->
-            print_error(
+            raise LivebookCLI.Error,
                   "Another application is already running on port #{port}." <>
                     " Either ensure this port is free or specify a different port using the --port option"
-            )
 
           :available ->
             start_server(extra_args)
@@ -108,7 +107,7 @@ defmodule LivebookCLI.Server do
         Process.sleep(:infinity)
 
       {:error, error} ->
-        print_error("Livebook failed to start with reason: #{inspect(error)}")
+        raise LivebookCLI.Error, "Livebook failed to start with reason: #{inspect(error)}"
     end
   end
 
@@ -126,20 +125,13 @@ defmodule LivebookCLI.Server do
 
   defp check_endpoint_availability(base_url) do
     Application.ensure_all_started(:req)
-
     health_url = set_path(base_url, "/public/health")
-
     req = Req.new() |> Livebook.Utils.req_attach_defaults()
 
     case Req.get(req, url: health_url, retry: false) do
-      {:ok, %{status: 200, body: %{"application" => "livebook"}}} ->
+      {:ok, %{status: 200, body: %{"application" => "livebook"}}} -> :livebook_running
-        :livebook_running
+      {:ok, _other} -> :taken
-
+      {:error, _exception} -> :available
-      {:ok, _other} ->
-        :taken
-
-      {:error, _exception} ->
-        :available
     end
   end
 
@@ -158,7 +150,7 @@ defmodule LivebookCLI.Server do
   end
 
   defp open_from_args(base_url, ["new"]) do
-    IO.warn(~s/passing "new" as an argument is deprecated, use "@new" instead/, [])
+    log_warning(~s/passing "new" as an argument is deprecated, use "@new" instead/)
     open_from_args(base_url, ["@new"])
   end
 
@@ -183,9 +175,9 @@ defmodule LivebookCLI.Server do
   end
 
   defp open_from_args(_base_url, _extra_args) do
-    print_error(
+    raise OptionParser.ParseError,
-      "Too many arguments entered. Ensure only one argument is used to specify the file path and all other arguments are preceded by the relevant switch"
+          "Too many arguments entered. Ensure only one argument is used to" <>
-    )
+            "specify the file path and all other arguments are preceded by the relevant switch"
   end
 
   @switches [
@@ -237,8 +229,4 @@ defmodule LivebookCLI.Server do
     |> Map.put(:path, path)
     |> URI.to_string()
   end
-
-  defp print_error(message) do
-    IO.ANSI.format([:red, message]) |> IO.puts()
-  end
 end

lib/livebook_cli/task.ex

@@ -1,11 +1,78 @@
 defmodule LivebookCLI.Task do
+  import LivebookCLI.Utils
+
   @doc """
   Returns a description of the task usage.
   """
-  @callback usage() :: String.t()
+  @callback usage() :: IO.chardata()
 
   @doc """
   Runs the task with the given list of command line arguments.
   """
   @callback call(args :: list(String.t())) :: :ok
+
+  @doc """
+  Runs the task with the given list of command line arguments.
+  """
+  @spec call(String.t(), list(String.t())) :: :ok
+  def call(name, args) do
+    task = fetch_task!(name)
+    task.call(args)
+  rescue
+    exception -> log_exception(exception, name, __STACKTRACE__)
+  end
+
+  @doc """
+  Shows the description of the task usage.
+  """
+  @spec usage(String.t()) :: :ok
+  def usage(name) do
+    task = fetch_task!(name)
+    log_info(task.usage())
+  rescue
+    exception -> log_exception(exception, name, __STACKTRACE__)
+  end
+
+  defp fetch_task!("server"), do: LivebookCLI.Server
+  defp fetch_task!("deploy"), do: LivebookCLI.Deploy
+
+  defp fetch_task!(name) do
+    log_error("Unknown command #{name}")
+    log_info(LivebookCLI.usage())
+
+    System.halt(1)
+  end
+
+  @spec log_exception(Exception.t(), String.t(), Exception.stacktrace()) :: no_return()
+  defp log_exception(exception, command_name, stacktrace) when is_exception(exception) do
+    [:red, format_exception(exception, command_name, stacktrace)]
+    |> IO.ANSI.format()
+    |> IO.puts()
+
+    System.halt(1)
+  end
+
+  defp format_exception(%OptionParser.ParseError{} = exception, command_name, _) do
+    """
+    #{Exception.message(exception)}
+
+    For more information try:
+
+       livebook #{command_name} --help
+    """
+  end
+
+  defp format_exception(%LivebookCLI.Error{} = exception, command_name, _) do
+    """
+    #{Exception.message(exception)}
+
+    For more information try:
+
+       livebook #{command_name} --help
+    """
+  end
+
+  defp format_exception(exception, _, stacktrace) do
+    Exception.format(:error, exception, stacktrace)
+  end
 end

lib/livebook_cli/utils.ex

@@ -0,0 +1,31 @@
+defmodule LivebookCLI.Utils do
+  def log_info(message) do
+    message
+    |> IO.ANSI.format()
+    |> IO.puts()
+  end
+
+  if Mix.env() == :dev do
+    def log_debug(message) do
+      [:cyan, message]
+      |> IO.ANSI.format()
+      |> IO.puts()
+    end
+  else
+    def log_debug(_message) do
+      :ok
+    end
+  end
+
+  def log_warning(message) do
+    [:yellow, message]
+    |> IO.ANSI.format()
+    |> IO.puts()
+  end
+
+  def log_error(message) do
+    [:red, message]
+    |> IO.ANSI.format()
+    |> IO.puts()
+  end
+end

test/livebook_teams/cli/deploy_test.exs

@@ -0,0 +1,271 @@
+defmodule LivebookCLI.Integration.DeployTest do
+  use Livebook.TeamsIntegrationCase, async: true
+
+  import Livebook.AppHelpers
+
+  alias Livebook.Utils
+
+  @url "http://localhost:4200"
+
+  @moduletag teams_for: :user
+  setup :teams
+
+  @moduletag subscribe_to_hubs_topics: [:connection]
+  @moduletag subscribe_to_teams_topics: [:clients, :deployment_groups, :app_deployments]
+
+  @moduletag :tmp_dir
+
+  describe "CLI deploy integration" do
+    test "successfully deploys a notebook via CLI",
+         %{team: team, node: node, org: org, tmp_dir: tmp_dir} do
+      title = "Test CLI Deploy App"
+      slug = Utils.random_short_id()
+      app_path = Path.join(tmp_dir, "#{slug}.livemd")
+      {key, _} = TeamsRPC.create_deploy_key(node, org: org)
+      deployment_group = TeamsRPC.create_deployment_group(node, org: org, url: @url)
+      hub_id = team.id
+      deployment_group_id = to_string(deployment_group.id)
+
+      stamp_notebook(app_path, """
+      <!-- livebook:{"app_settings":{"access_type":"public","slug":"#{slug}"},"hub_id":"#{hub_id}"} -->
+
+      # #{title}
+
+      ## Test Section
+
+      ```elixir
+      IO.puts("Hello from CLI deployed app!")
+      ```
+      """)
+
+      output =
+        ExUnit.CaptureIO.capture_io(fn ->
+          assert deploy(
+                   key,
+                   team.teams_key,
+                   deployment_group.name,
+                   app_path
+                 ) == :ok
+        end)
+
+      assert output =~ "* Preparing to deploy notebook #{slug}.livemd"
+      assert output =~ "  * #{title} deployed successfully. (#{@url}/apps/#{slug})"
+
+      assert_receive {:app_deployment_started,
+                      %{
+                        title: ^title,
+                        slug: ^slug,
+                        deployment_group_id: ^deployment_group_id,
+                        hub_id: ^hub_id,
+                        deployed_by: "CLI"
+                      }}
+    end
+
+    test "successfully deploys multiple notebooks from directory",
+         %{team: team, node: node, org: org, tmp_dir: tmp_dir} do
+      {key, _} = TeamsRPC.create_deploy_key(node, org: org)
+      deployment_group = TeamsRPC.create_deployment_group(node, org: org, url: @url)
+      hub_id = team.id
+      deployment_group_id = to_string(deployment_group.id)
+
+      apps =
+        for i <- 1..3 do
+          title = "Test App #{i}"
+          slug = "app-#{i}-#{Utils.random_short_id()}"
+          app_path = Path.join(tmp_dir, "#{slug}.livemd")
+
+          stamp_notebook(app_path, """
+          <!-- livebook:{"app_settings":{"access_type":"public","slug":"#{slug}"},"hub_id":"#{hub_id}"} -->
+
+          # #{title}
+
+          ```elixir
+          IO.puts("Hello from app #{i}!")
+          ```
+          """)
+
+          {slug, title}
+        end
+
+      output =
+        ExUnit.CaptureIO.capture_io(fn ->
+          assert deploy(
+                   key,
+                   team.teams_key,
+                   deployment_group.name,
+                   Path.join(tmp_dir, "*.livemd")
+                 ) == :ok
+        end)
+
+      for {slug, title} <- apps do
+        assert output =~ "* Preparing to deploy notebook #{slug}.livemd"
+        assert output =~ "  * #{title} deployed successfully. (#{@url}/apps/#{slug})"
+
+        assert_receive {:app_deployment_started,
+                        %{
+                          title: ^title,
+                          slug: ^slug,
+                          deployment_group_id: ^deployment_group_id,
+                          hub_id: ^hub_id,
+                          deployed_by: "CLI"
+                        }}
+      end
+    end
+
+    test "fails with invalid deploy key", %{team: team, node: node, org: org, tmp_dir: tmp_dir} do
+      slug = Utils.random_short_id()
+      app_path = Path.join(tmp_dir, "#{slug}.livemd")
+      deployment_group = TeamsRPC.create_deployment_group(node, org: org, url: @url)
+
+      stamp_notebook(app_path, """
+      <!-- livebook:{"app_settings":{"access_type":"public","slug":"#{slug}"},"hub_id":"#{team.id}"} -->
+
+      # Test App
+      """)
+
+      assert_raise LivebookCLI.Error, ~r/Deploy Key must be a Livebook Teams Deploy Key/s, fn ->
+        deploy(
+          "invalid_key",
+          team.teams_key,
+          deployment_group.name,
+          app_path
+        )
+      end
+    end
+
+    test "fails with invalid teams key", %{team: team, node: node, org: org, tmp_dir: tmp_dir} do
+      slug = Utils.random_short_id()
+      app_path = Path.join(tmp_dir, "#{slug}.livemd")
+      {key, _} = TeamsRPC.create_deploy_key(node, org: org)
+      deployment_group = TeamsRPC.create_deployment_group(node, org: org, url: @url)
+
+      stamp_notebook(app_path, """
+      <!-- livebook:{"app_settings":{"access_type":"public","slug":"#{slug}"},"hub_id":"#{team.id}"} -->
+
+      # Test App
+      """)
+
+      assert_raise LivebookCLI.Error, ~r/Teams Key must be a Livebook Teams Key/s, fn ->
+        deploy(
+          key,
+          "invalid-key",
+          deployment_group.name,
+          app_path
+        )
+      end
+    end
+
+    test "fails with missing deployment group",
+         %{team: team, node: node, org: org, tmp_dir: tmp_dir} do
+      slug = Utils.random_short_id()
+      app_path = Path.join(tmp_dir, "#{slug}.livemd")
+      {key, _} = TeamsRPC.create_deploy_key(node, org: org)
+
+      stamp_notebook(app_path, """
+      <!-- livebook:{"app_settings":{"access_type":"public","slug":"#{slug}"},"hub_id":"#{team.id}"} -->
+
+      # Test App
+      """)
+
+      assert_raise LivebookCLI.Error, ~r/Deployment Group can't be blank/s, fn ->
+        deploy(
+          key,
+          team.teams_key,
+          "",
+          app_path
+        )
+      end
+    end
+
+    test "fails with invalid deployment group",
+         %{team: team, node: node, org: org, tmp_dir: tmp_dir} do
+      title = "Test CLI Deploy App"
+      slug = Utils.random_short_id()
+      app_path = Path.join(tmp_dir, "#{slug}.livemd")
+      {key, _} = TeamsRPC.create_deploy_key(node, org: org)
+      deployment_group = TeamsRPC.create_deployment_group(node, org: org, url: @url)
+      hub_id = team.id
+      deployment_group_id = to_string(deployment_group.id)
+
+      stamp_notebook(app_path, """
+      <!-- livebook:{"app_settings":{"access_type":"public","slug":"#{slug}"},"hub_id":"#{hub_id}"} -->
+
+      # #{title}
+
+      ## Test Section
+
+      ```elixir
+      IO.puts("Hello from CLI deployed app!")
+      ```
+      """)
+
+      assert_raise LivebookCLI.Error, ~r/Deployment Group does not exist/s, fn ->
+        ExUnit.CaptureIO.capture_io(fn ->
+          deploy(
+            key,
+            team.teams_key,
+            Utils.random_short_id(),
+            app_path
+          )
+        end)
+      end
+
+      refute_receive {:app_deployment_started,
+                      %{
+                        title: ^title,
+                        slug: ^slug,
+                        deployment_group_id: ^deployment_group_id,
+                        hub_id: ^hub_id,
+                        deployed_by: "CLI"
+                      }}
+    end
+
+    test "fails with non-existent file", %{team: team, node: node, org: org, tmp_dir: tmp_dir} do
+      {key, _} = TeamsRPC.create_deploy_key(node, org: org)
+      deployment_group = TeamsRPC.create_deployment_group(node, org: org, url: @url)
+
+      assert_raise LivebookCLI.Error, ~r/File Paths must be a valid path/s, fn ->
+        deploy(
+          key,
+          team.teams_key,
+          deployment_group.name,
+          Path.join(tmp_dir, "app.livemd")
+        )
+      end
+    end
+
+    test "fails with directory argument", %{team: team, node: node, org: org, tmp_dir: tmp_dir} do
+      {key, _} = TeamsRPC.create_deploy_key(node, org: org)
+      deployment_group = TeamsRPC.create_deployment_group(node, org: org, url: @url)
+
+      assert_raise LivebookCLI.Error, ~r/File Paths must be a file path/s, fn ->
+        deploy(
+          key,
+          team.teams_key,
+          deployment_group.name,
+          tmp_dir
+        )
+      end
+    end
+  end
+
+  defp deploy(deploy_key, teams_key, deployment_group_name, path) do
+    paths =
+      case Path.wildcard(path) do
+        [] -> [path]
+        [path] -> [path]
+        paths -> paths
+      end
+
+    LivebookCLI.Deploy.call(
+      [
+        "--deploy-key",
+        deploy_key,
+        "--teams-key",
+        teams_key,
+        "--deployment-group",
+        deployment_group_name
+      ] ++ paths
+    )
+  end
+end

test/livebook_teams/teams_test.exs

@@ -250,4 +250,128 @@ defmodule Livebook.TeamsTest do
       assert_receive {:app_deployment_stopped, ^app_deployment2}
     end
   end
+
+  describe "fetch_cli_session/1" do
+    @describetag teams_for: :cli
+
+    @tag teams_persisted: false
+    test "authenticates the deploy key", %{team: team} do
+      config = %{teams_key: team.teams_key, session_token: team.session_token}
+
+      refute Livebook.Hubs.hub_exists?(team.id)
+      assert Teams.fetch_cli_session(config) == {:ok, team}
+      assert Livebook.Hubs.hub_exists?(team.id)
+    end
+
+    @tag teams_for: :user
+    test "authenticates the deploy key when hub already exists",
+         %{team: team, org: org, node: node} do
+      {key, _} = TeamsRPC.create_deploy_key(node, org: org)
+      config = %{teams_key: team.teams_key, session_token: key}
+
+      assert Teams.fetch_cli_session(config) ==
+               {:ok,
+                %Livebook.Hubs.Team{
+                  billing_status: team.billing_status,
+                  hub_emoji: "🚀",
+                  hub_name: team.hub_name,
+                  id: team.id,
+                  offline: nil,
+                  org_id: team.org_id,
+                  org_key_id: team.org_key_id,
+                  org_public_key: team.org_public_key,
+                  session_token: key,
+                  teams_key: team.teams_key,
+                  user_id: nil
+                }}
+
+      assert Livebook.Hubs.hub_exists?(team.id)
+      assert Livebook.Hubs.fetch_hub!(team.id).session_token == key
+    end
+
+    @tag teams_persisted: false
+    test "returns error with invalid credentials", %{team: team} do
+      config = %{teams_key: team.teams_key, session_token: "lb_dk_foo"}
+
+      assert {:transport_error, "You are not authorized" <> _} = Teams.fetch_cli_session(config)
+      refute Livebook.Hubs.hub_exists?(team.id)
+    end
+  end
+
+  describe "deploy_app_from_cli/2" do
+    @describetag teams_for: :user
+
+    @tag :tmp_dir
+    test "deploys app to Teams using a CLI session",
+         %{team: team, node: node, tmp_dir: tmp_dir, org: org} do
+      %{id: id, name: name} =
+        TeamsRPC.create_deployment_group(node,
+          name: "angry-cat-#{Ecto.UUID.generate()}",
+          url: "http://localhost:4123",
+          mode: :online,
+          org: org
+        )
+
+      id = to_string(id)
+      hub_id = "team-#{org.name}"
+      slug = Utils.random_short_id()
+      title = "MyNotebook-#{slug}"
+      app_settings = %{Notebook.AppSettings.new() | slug: slug}
+
+      notebook = %{
+        Notebook.new()
+        | app_settings: app_settings,
+          name: title,
+          hub_id: hub_id,
+          deployment_group_id: id
+      }
+
+      files_dir = FileSystem.File.local(tmp_dir)
+
+      # stamp the notebook
+      assert {:ok, app_deployment} = Teams.AppDeployment.new(notebook, files_dir)
+
+      # fetch the cli session
+      {key, _deploy_key} = TeamsRPC.create_deploy_key(node, org: org)
+      config = %{teams_key: team.teams_key, session_token: key}
+      assert {:ok, team} = Teams.fetch_cli_session(config)
+
+      # deploy the app
+      assert {:ok, _url} = Teams.deploy_app_from_cli(team, app_deployment, name)
+
+      sha = app_deployment.sha
+      multi_session = app_settings.multi_session
+      access_type = app_settings.access_type
+
+      assert_receive {:app_deployment_started,
+                      %Livebook.Teams.AppDeployment{
+                        slug: ^slug,
+                        sha: ^sha,
+                        title: ^title,
+                        deployed_by: "CLI",
+                        multi_session: ^multi_session,
+                        access_type: ^access_type,
+                        deployment_group_id: ^id
+                      } = app_deployment2}
+
+      assert Teams.deploy_app_from_cli(team, app_deployment, "foo") ==
+               {:error, %{"deployment_group" => ["does not exist"]}}
+
+      assert Teams.deploy_app_from_cli(team, %{app_deployment | slug: "@abc"}, name) ==
+               {:error, %{"slug" => ["should only contain alphanumeric characters and dashes"]}}
+
+      assert Teams.deploy_app_from_cli(team, %{app_deployment | multi_session: nil}, name) ==
+               {:error, %{"multi_session" => ["can't be blank"]}}
+
+      assert Teams.deploy_app_from_cli(team, %{app_deployment | access_type: nil}, name) ==
+               {:error, %{"access_type" => ["can't be blank"]}}
+
+      assert Teams.deploy_app_from_cli(team, %{app_deployment | access_type: :abc}, name) ==
+               {:error, %{"access_type" => ["is invalid"]}}
+
+      # force app deployment to be stopped
+      TeamsRPC.toggle_app_deployment(node, app_deployment2.id, team.org_id)
+      assert_receive {:app_deployment_stopped, ^app_deployment2}
+    end
+  end
 end

test/support/factory.ex

@@ -21,7 +21,8 @@ defmodule Livebook.Factory do
       org_id: 1,
       user_id: 1,
       org_key_id: 1,
-      org_public_key: Livebook.Hubs.Team.public_key_prefix() <> Livebook.Utils.random_long_id(),
+      org_public_key:
+        Livebook.Teams.Constants.public_key_prefix() <> Livebook.Utils.random_long_id(),
       teams_key: org.teams_key,
       session_token: Livebook.Utils.random_short_id(),
       offline: nil

test/support/integration/teams_rpc.ex

@@ -137,6 +137,11 @@ defmodule Livebook.TeamsRPC do
     :erpc.call(node, TeamsRPC, :create_authorization_group, [attrs])
   end
 
+  def create_deploy_key(node, attrs \\ []) do
+    key = :erpc.call(node, TeamsRPC, :generate_deploy_key, [])
+    {key, :erpc.call(node, TeamsRPC, :create_deploy_key, [key, attrs])}
+  end
+
   # Update resource
 
   def update_authorization_group(node, authorization_group, attrs) do

test/support/integration/teams_tests.ex

@@ -12,13 +12,15 @@ defmodule Livebook.TeamsIntegrationHelper do
       {:user, _} -> Map.merge(context, create_user_hub(context.node))
       {:agent, false} -> Map.merge(context, new_agent_hub(context.node))
       {:agent, _} -> Map.merge(context, create_agent_hub(context.node))
+      {:cli, false} -> Map.merge(context, new_cli_hub(context.node))
+      {:cli, _} -> Map.merge(context, create_cli_hub(context.node))
       _otherwise -> context
     end
   end
 
-  def livebook_teams_auth(%{conn: conn, node: node, team: team} = context) do
+  def livebook_teams_auth(%{node: node, team: team} = context) do
     ZTA.LivebookTeams.start_link(name: context.test, identity_key: team.id)
-    {conn, code} = authenticate_user_on_teams(context.test, conn, node, team)
+    {conn, code} = authenticate_user_on_teams(context.test, node, team)
 
     Map.merge(context, %{conn: conn, code: code})
   end
@@ -123,12 +125,64 @@ defmodule Livebook.TeamsIntegrationHelper do
     }
   end
 
-  defp authenticate_user_on_teams(name, conn, node, team) do
+  def create_cli_hub(node, opts \\ []) do
-    # Create a fresh connection to avoid session contamination
+    context = new_cli_hub(node, opts)
-    fresh_conn = Phoenix.ConnTest.build_conn()
+
+    Hubs.save_hub(context.team)
+    ExUnit.Callbacks.on_exit(fn -> Hubs.delete_hub(context.team.id) end)
+
+    %{context | team: Hubs.fetch_hub!(context.team.id)}
+  end
+
+  def new_cli_hub(node, opts \\ []) do
+    {teams_key, key_hash} = generate_key_hash()
+
+    org = TeamsRPC.create_org(node)
+    org_key = TeamsRPC.create_org_key(node, org: org, key_hash: key_hash)
+    org_key_pair = TeamsRPC.create_org_key_pair(node, org: org)
+
+    attrs =
+      opts
+      |> Keyword.get(:deployment_group, [])
+      |> Keyword.merge(
+        name: "angry-cat-#{Ecto.UUID.generate()}",
+        mode: :online,
+        org: org
+      )
+
+    deployment_group = TeamsRPC.create_deployment_group(node, attrs)
+    {key, deploy_key} = TeamsRPC.create_deploy_key(node, org: org)
+
+    TeamsRPC.create_billing_subscription(node, org)
+
+    team =
+      Factory.build(:team,
+        id: "team-#{org.name}",
+        hub_name: org.name,
+        hub_emoji: "🚀",
+        user_id: nil,
+        org_id: org.id,
+        org_key_id: org_key.id,
+        org_public_key: org_key_pair.public_key,
+        session_token: key,
+        teams_key: teams_key
+      )
+
+    %{
+      deploy_key: Map.replace!(deploy_key, :key_hash, key),
+      deployment_group: deployment_group,
+      org: org,
+      org_key: org_key,
+      org_key_pair: org_key_pair,
+      team: team
+    }
+  end
+
+  def authenticate_user_on_teams(name, node, team) do
+    conn = Phoenix.ConnTest.build_conn()
 
     response =
-      fresh_conn
+      conn
       |> LivebookWeb.ConnCase.with_authorization(team.id, name)
       |> get("/")
       |> html_response(200)
@@ -140,12 +194,11 @@ defmodule Livebook.TeamsIntegrationHelper do
     %{code: code} = Livebook.TeamsRPC.allow_auth_request(node, token)
 
     session =
-      fresh_conn
+      conn
       |> LivebookWeb.ConnCase.with_authorization(team.id, name)
       |> get("/", %{teams_identity: "", code: code})
       |> Plug.Conn.get_session()
 
-    # Initialize the original conn with the new session data
     authenticated_conn = Plug.Test.init_test_session(conn, session)
     final_conn = get(authenticated_conn, "/")
     assigns = Map.take(final_conn.assigns, [:current_user])