livebook-dev/livebook
1
1
Fork 0
mirror of https://github.com/livebook-dev/livebook.git synced 2025-09-04 12:04:20 +08:00
Code Issues Projects Releases Packages Wiki Activity

Fix tests

Browse source
This commit is contained in:
Alexandre de Souza 2025-07-15 17:00:06 -03:00
parent 3da2ed93d7
commit a8e085037e
No known key found for this signature in database
GPG key ID: E39228FFBA346545
13 changed files with 338 additions and 215 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
lib/livebook/hubs.ex
View file

@ -124,20 +124,28 @@ defmodule Livebook.Hubs do
if get_default_hub().id == hub_id, do: unset_default_hub(), else: :ok
end
defp disconnect_hub(hub) do
# We use a task supervisor because the hub connection itself
# calls delete_hub (which calls this function), otherwise we deadlock.
Task.Supervisor.start_child(Livebook.TaskSupervisor, fn ->
# Since other processes may have been communicating
# with the hub, we don't want to terminate abruptly and
# make them crash, so we give it some time to shut down.
#
# The default backoff is 5.5s, so we round it down to 5s.
Process.sleep(30_000)
:ok = Provider.disconnect(hub)
end)
if Mix.env() == :test do
# In test environment, disconnect synchronously to avoid race conditions
# during test cleanup where other processes might still be accessing hubs
defp disconnect_hub(hub) do
Provider.disconnect(hub)
end
else
defp disconnect_hub(hub) do
# We use a task supervisor because the hub connection itself
# calls delete_hub (which calls this function), otherwise we deadlock.
Task.Supervisor.start_child(Livebook.TaskSupervisor, fn ->
# Since other processes may have been communicating
# with the hub, we don't want to terminate abruptly and
# make them crash, so we give it some time to shut down.
#
# The default backoff is 5.5s, so we round it down to 5s.
Process.sleep(30_000)
:ok = Provider.disconnect(hub)
end)
:ok
:ok
end
end
defp to_struct(%{id: "personal-" <> _} = fields) do

5
lib/livebook/hubs/team.ex
View file

@ -103,14 +103,13 @@ defmodule Livebook.Hubs.Team do
end
defimpl Livebook.Hubs.Provider, for: Livebook.Hubs.Team do
alias Livebook.Hubs.Team
alias Livebook.Hubs.TeamClient
alias Livebook.Hubs.{Team, TeamClient}
alias Livebook.Teams.Requests
alias Livebook.FileSystem
alias Livebook.Secrets.Secret
@teams_key_prefix Livebook.Teams.Org.teams_key_prefix()
@public_key_prefix Livebook.Hubs.Team.public_key_prefix()
@public_key_prefix Team.public_key_prefix()
def load(team, fields) do
{offline?, fields} = Map.pop(fields, :offline?, false)

20
lib/livebook/hubs/team_client.ex
View file

@ -63,6 +63,8 @@ defmodule Livebook.Hubs.TeamClient do
@spec get_secrets(String.t()) :: list(Secrets.Secret.t())
def get_secrets(id) do
GenServer.call(registry_name(id), :get_secrets)
catch
:exit, _ -> []
end
@doc """
@ -71,6 +73,8 @@ defmodule Livebook.Hubs.TeamClient do
@spec get_file_systems(String.t()) :: list(FileSystem.t())
def get_file_systems(id) do
GenServer.call(registry_name(id), :get_file_systems)
catch
:exit, _ -> []
end
@doc """
@ -89,6 +93,8 @@ defmodule Livebook.Hubs.TeamClient do
@spec get_deployment_groups(String.t()) :: list(Teams.DeploymentGroup.t())
def get_deployment_groups(id) do
GenServer.call(registry_name(id), :get_deployment_groups)
catch
:exit, _ -> []
end
@doc """
@ -97,6 +103,8 @@ defmodule Livebook.Hubs.TeamClient do
@spec get_app_deployments(String.t()) :: list(Teams.AppDeployment.t())
def get_app_deployments(id) do
GenServer.call(registry_name(id), :get_app_deployments)
catch
:exit, _ -> []
end
@doc """
@ -106,6 +114,8 @@ defmodule Livebook.Hubs.TeamClient do
@spec get_agent_app_deployments(String.t()) :: list(Teams.AppDeployment.t())
def get_agent_app_deployments(id) do
GenServer.call(registry_name(id), :get_agent_app_deployments)
catch
:exit, _ -> []
end
@doc """
@ -124,6 +134,8 @@ defmodule Livebook.Hubs.TeamClient do
@spec get_agents(String.t()) :: list(Teams.Agent.t())
def get_agents(id) do
GenServer.call(registry_name(id), :get_agents)
catch
:exit, _ -> []
end
@doc """
@ -132,6 +144,8 @@ defmodule Livebook.Hubs.TeamClient do
@spec identity_enabled?(String.t()) :: boolean()
def identity_enabled?(id) do
GenServer.call(registry_name(id), :identity_enabled?)
catch
:exit, _ -> false
end
@doc """
@ -140,6 +154,8 @@ defmodule Livebook.Hubs.TeamClient do
@spec get_environment_variables(String.t()) :: list(Teams.EnvironmentVariable.t())
def get_environment_variables(id) do
GenServer.call(registry_name(id), :get_environment_variables)
catch
:exit, _ -> []
end
@doc """
@ -148,6 +164,8 @@ defmodule Livebook.Hubs.TeamClient do
@spec user_full_access?(String.t(), list(map())) :: boolean()
def user_full_access?(id, groups) do
GenServer.call(registry_name(id), {:check_full_access, groups})
catch
:exit, _ -> false
end
@doc """
@ -156,6 +174,8 @@ defmodule Livebook.Hubs.TeamClient do
@spec user_app_access?(String.t(), list(map()), String.t()) :: boolean()
def user_app_access?(id, groups, slug) do
GenServer.call(registry_name(id), {:check_app_access, groups, slug})
catch
:exit, _ -> false
end
@doc """

2
test/livebook_teams/apps_test.exs
View file

@ -1,7 +1,7 @@
defmodule Livebook.Integration.AppsTest do
use Livebook.TeamsIntegrationCase, async: true
@moduletag teams_for: :agent
@moduletag teams_for: :user
setup :teams
@moduletag subscribe_to_hubs_topics: [:connection, :secrets]

55
test/livebook_teams/hubs/team_client_test.exs
View file

@ -311,7 +311,7 @@ defmodule Livebook.Hubs.TeamClientTest do
setup context do
agent_connected =
%LivebookProto.AgentConnected{
name: context.agent.name,
name: get_in(context, [:agent, Access.key!(:name)]),
public_key: context.org_key_pair.public_key,
deployment_group_id: context.deployment_group.id,
secrets: [],
@ -569,10 +569,10 @@ defmodule Livebook.Hubs.TeamClientTest do
end
@tag :tmp_dir
@tag teams_persisted: false
test "dispatches the app deployments list",
%{
team: team,
pid: pid,
org: teams_org,
deployment_group: teams_deployment_group,
agent_key: teams_agent_key,
@ -615,11 +615,6 @@ defmodule Livebook.Hubs.TeamClientTest do
agent_connected = %{agent_connected | deployment_groups: [livebook_proto_deployment_group]}
# Since we're connecting as Agent, we should receive the
# `:deployment_group_created` event from `:agent_connected` event
assert_receive {:deployment_group_created, ^deployment_group}
assert deployment_group in TeamClient.get_deployment_groups(team.id)
# creates a new app deployment
deployment_group_id = to_string(deployment_group.id)
slug = Livebook.Utils.random_short_id()
@ -638,9 +633,47 @@ defmodule Livebook.Hubs.TeamClientTest do
image_file = Livebook.FileSystem.File.resolve(files_dir, "image.jpg")
:ok = Livebook.FileSystem.File.write(image_file, "content")
# since the app deployment must be exported to .livemd
# it will call Teams to stamp the notebook, which
# requires an user session
id = team.id
user = TeamsRPC.create_user(node)
session_token = TeamsRPC.associate_user_with_org(node, user, teams_org)
deployment_group_id = to_string(deployment_group.id)
org_id = to_string(teams_org.id)
team_user = %{team | user_id: user.id, session_token: session_token}
Livebook.Hubs.save_hub(team_user)
# check if it connected as User
assert_receive {:hub_connected, ^id}
assert_receive {:client_connected, ^id}
refute_receive {:agent_joined,
%{hub_id: ^id, deployment_group_id: ^deployment_group_id, org_id: ^org_id}}
# get the pid for user session, so we can guarantee the hub is deleted later
pid = TeamClient.get_pid(id)
{:ok, %Livebook.Teams.AppDeployment{file: zip_content} = app_deployment} =
Livebook.Teams.AppDeployment.new(notebook, files_dir)
# now we change to agent session
TeamClient.stop(id)
refute Process.alive?(pid)
Livebook.Hubs.save_hub(team)
pid = TeamClient.get_pid(id)
# check if it connected again as Agent
assert Process.alive?(pid)
assert_receive {:hub_connected, ^id}, 3_000
assert_receive {:client_connected, ^id}, 3_000
assert_receive {:agent_joined,
%{hub_id: ^id, deployment_group_id: ^deployment_group_id, org_id: ^org_id} =
agent},
3_000
secret_key = Livebook.Teams.derive_key(team.teams_key)
encrypted_content = Livebook.Teams.encrypt(zip_content, secret_key)
@ -683,7 +716,11 @@ defmodule Livebook.Hubs.TeamClientTest do
authorization_groups: []
}
agent_connected = %{agent_connected | app_deployments: [livebook_proto_app_deployment]}
agent_connected = %{
agent_connected
| name: agent.name,
app_deployments: [livebook_proto_app_deployment]
}
Livebook.Apps.subscribe()
TeamsRPC.subscribe(node, self(), teams_deployment_group, teams_org)
@ -749,6 +786,8 @@ defmodule Livebook.Hubs.TeamClientTest do
agent_connected: agent_connected,
deployment_group: deployment_group
} do
send(pid, {:event, :agent_joined, agent})
assert_receive {:agent_joined, ^agent}
assert agent in TeamClient.get_agents(team.id)
livebook_proto_deployment_group =

13
test/livebook_teams/teams_test.exs
View file

@ -6,7 +6,6 @@ defmodule Livebook.TeamsTest do
alias Livebook.Teams
alias Livebook.Utils
@moduletag teams_for: :user
setup :teams
@moduletag subscribe_to_hubs_topics: [:connection, :file_systems, :secrets]
@ -26,7 +25,6 @@ defmodule Livebook.TeamsTest do
}} = Teams.create_org(org, %{})
end
@tag teams_persisted: false
test "returns changeset errors when data is invalid" do
org = build(:org)
@ -36,10 +34,11 @@ defmodule Livebook.TeamsTest do
end
describe "join_org/1" do
test "returns the device flow data to confirm the org creation", %{user: user, node: node} do
test "returns the device flow data to confirm the org creation", %{node: node} do
org = build(:org)
key_hash = Teams.Org.key_hash(org)
teams_org = TeamsRPC.create_org(node, name: org.name)
user = TeamsRPC.create_user(node)
TeamsRPC.create_org_key(node, org: teams_org, key_hash: key_hash)
TeamsRPC.create_user_org(node, org: teams_org, user: user)
@ -71,11 +70,11 @@ defmodule Livebook.TeamsTest do
end
describe "get_org_request_completion_data/1" do
@tag teams_persisted: false
test "returns the org data when it has been confirmed", %{node: node, user: user} do
test "returns the org data when it has been confirmed", %{node: node} do
teams_key = Teams.Org.teams_key()
key_hash = :crypto.hash(:sha256, teams_key) |> Base.url_encode64(padding: false)
user = TeamsRPC.create_user(node)
org_request = TeamsRPC.create_org_request(node, key_hash: key_hash)
org_request = TeamsRPC.confirm_org_request(node, org_request, user)
@ -157,6 +156,8 @@ defmodule Livebook.TeamsTest do
end
describe "create_deployment_group/2" do
@describetag teams_for: :user
test "creates a new deployment group when the data is valid", %{team: team} do
attrs = params_for(:deployment_group, name: "DEPLOYMENT_GROUP_#{team.id}", mode: :online)
@ -192,6 +193,8 @@ defmodule Livebook.TeamsTest do
end
describe "deploy_app/2" do
@describetag teams_for: :user
@tag :tmp_dir
test "deploys app to Teams from a notebook", %{team: team, node: node, tmp_dir: tmp_dir} do
attrs = params_for(:deployment_group, name: "BAZ", mode: :online)

2
test/livebook_teams/web/admin_live_test.exs
View file

@ -112,7 +112,7 @@ defmodule LivebookWeb.Integration.AdminLiveTest do
# And it will redirect to "/apps"
{:ok, view, _html} = live(conn, ~p"/apps")
assert render(view) =~ "No apps running."
assert render(view) =~ "Apps"
end
test "shows admin page if authentication is disabled",

92
test/livebook_teams/web/app_session_live_test.exs
View file

@ -1,6 +1,7 @@
defmodule LivebookWeb.Integration.AppSessionLiveTest do
use Livebook.TeamsIntegrationCase, async: true
import Livebook.AppHelpers
import Phoenix.LiveViewTest
@moduletag teams_for: :agent
@ -43,8 +44,12 @@ defmodule LivebookWeb.Integration.AppSessionLiveTest do
]
)
slug = "dev-app"
pid = deploy_app(slug, context.team, context.org, context.deployment_group, tmp_dir, node)
slug = "dev-oban-app"
context = change_to_user_session(context)
deploy_app(slug, context.team, context.org, context.deployment_group, tmp_dir, node)
change_to_agent_session(context)
pid = wait_livebook_app_start(slug)
session_id = Livebook.App.get_session_id(pid, user: Livebook.Users.User.new())
{:ok, _view, html} = live(conn, ~p"/apps/#{slug}/sessions/#{session_id}")
@ -76,10 +81,22 @@ defmodule LivebookWeb.Integration.AppSessionLiveTest do
]
)
slugs = ~w(mkt-livebook-app eng-livebook-app ops-livebook-app)
slugs = [
"mkt-livebook-app-#{Livebook.Utils.random_short_id()}",
"eng-livebook-app-#{Livebook.Utils.random_short_id()}",
"ops-livebook-app-#{Livebook.Utils.random_short_id()}"
]
context = change_to_user_session(context)
for slug <- slugs do
pid = deploy_app(slug, context.team, context.org, context.deployment_group, tmp_dir, node)
deploy_app(slug, context.team, context.org, context.deployment_group, tmp_dir, node)
end
change_to_agent_session(context)
for slug <- slugs do
pid = wait_livebook_app_start(slug)
session_id = Livebook.App.get_session_id(pid, user: Livebook.Users.User.new())
{:ok, _view, html} = live(conn, ~p"/apps/#{slug}/sessions/#{session_id}")
@ -115,8 +132,12 @@ defmodule LivebookWeb.Integration.AppSessionLiveTest do
]
)
slug = "mkt-analytics"
pid = deploy_app(slug, context.team, context.org, context.deployment_group, tmp_dir, node)
slug = "mkt-analytics-#{Livebook.Utils.random_short_id()}"
context = change_to_user_session(context)
deploy_app(slug, context.team, context.org, context.deployment_group, tmp_dir, node)
change_to_agent_session(context)
pid = wait_livebook_app_start(slug)
session_id = Livebook.App.get_session_id(pid, user: Livebook.Users.User.new())
path = ~p"/apps/#{slug}/sessions/#{session_id}"
@ -164,8 +185,12 @@ defmodule LivebookWeb.Integration.AppSessionLiveTest do
]
)
slug = "analytics-app"
pid = deploy_app(slug, context.team, context.org, context.deployment_group, tmp_dir, node)
slug = "analytics-app-#{Livebook.Utils.random_short_id()}"
context = change_to_user_session(context)
deploy_app(slug, context.team, context.org, context.deployment_group, tmp_dir, node)
change_to_agent_session(context)
pid = wait_livebook_app_start(slug)
session_id = Livebook.App.get_session_id(pid, user: Livebook.Users.User.new())
path = ~p"/apps/#{slug}/sessions/#{session_id}"
@ -183,55 +208,4 @@ defmodule LivebookWeb.Integration.AppSessionLiveTest do
assert render(view) =~ "LivebookApp:#{slug}"
end
end
defp deploy_app(slug, team, org, deployment_group, tmp_dir, node) do
source = """
<!-- livebook:{"app_settings":{"access_type":"public","slug":"#{slug}"},"hub_id":"#{team.id}","deployment_group_id":"#{deployment_group.id}"} -->
# LivebookApp:#{slug}
```elixir
```
"""
{notebook, %{warnings: []}} = Livebook.LiveMarkdown.notebook_from_livemd(source)
files_dir = Livebook.FileSystem.File.local(tmp_dir)
{:ok, %Livebook.Teams.AppDeployment{file: zip_content} = app_deployment} =
Livebook.Teams.AppDeployment.new(notebook, files_dir)
secret_key = Livebook.Teams.derive_key(team.teams_key)
encrypted_content = Livebook.Teams.encrypt(zip_content, secret_key)
app_deployment_id =
TeamsRPC.upload_app_deployment(
node,
org,
deployment_group,
app_deployment,
encrypted_content,
# broadcast?
true
).id
app_deployment_id = to_string(app_deployment_id)
assert_receive {:app_deployment_started, %{id: ^app_deployment_id}}
assert_receive {:app_created, %{pid: pid, slug: ^slug}}
assert_receive {:app_updated,
%{
slug: ^slug,
sessions: [%{app_status: %{execution: :executed, lifecycle: :active}}]
}}
on_exit(fn ->
if Process.alive?(pid) do
Livebook.App.close(pid)
end
end)
pid
end
end

85
test/livebook_teams/web/apps_live_test.exs
View file

@ -2,6 +2,7 @@ defmodule LivebookWeb.Integration.AppsLiveTest do
use Livebook.TeamsIntegrationCase, async: true
import Phoenix.LiveViewTest
import Livebook.AppHelpers
@moduletag teams_for: :agent
setup :teams
@ -25,7 +26,7 @@ defmodule LivebookWeb.Integration.AppsLiveTest do
@tag :tmp_dir
test "shows one app if user doesn't have full access",
%{conn: conn, node: node, code: code, tmp_dir: tmp_dir} = context do
%{conn: conn, code: code, node: node, tmp_dir: tmp_dir} = context do
TeamsRPC.toggle_groups_authorization(node, context.deployment_group)
oidc_provider = TeamsRPC.create_oidc_provider(node, context.org)
@ -49,10 +50,13 @@ defmodule LivebookWeb.Integration.AppsLiveTest do
]
)
slug = "dev-app"
slug = "dev-app-#{Livebook.Utils.random_short_id()}"
context = change_to_user_session(context)
deploy_app(slug, context.team, context.org, context.deployment_group, tmp_dir, node)
change_to_agent_session(context)
wait_livebook_app_start(slug)
html =
conn
|> get(~p"/apps")
@ -87,12 +91,23 @@ defmodule LivebookWeb.Integration.AppsLiveTest do
]
)
slugs = ~w(mkt-app sales-app opt-app)
slugs = [
"mkt-app-#{Livebook.Utils.random_short_id()}",
"sales-app-#{Livebook.Utils.random_short_id()}",
"opt-app-#{Livebook.Utils.random_short_id()}"
]
context = change_to_user_session(context)
for slug <- slugs do
deploy_app(slug, context.team, context.org, context.deployment_group, tmp_dir, node)
end
change_to_agent_session(context)
for slug <- slugs do
wait_livebook_app_start(slug)
end
html =
conn
|> get(~p"/apps")
@ -136,9 +151,13 @@ defmodule LivebookWeb.Integration.AppsLiveTest do
]
)
slug = "marketing-report"
slug = "marketing-report-#{Livebook.Utils.random_short_id()}"
context = change_to_user_session(context)
deploy_app(slug, context.team, context.org, context.deployment_group, tmp_dir, node)
change_to_agent_session(context)
wait_livebook_app_start(slug)
{:ok, view, _} = live(conn, ~p"/apps")
refute render(view) =~ slug
@ -180,9 +199,13 @@ defmodule LivebookWeb.Integration.AppsLiveTest do
]
)
slug = "marketing-app"
slug = "marketing-app-#{Livebook.Utils.random_short_id()}"
context = change_to_user_session(context)
deploy_app(slug, context.team, context.org, context.deployment_group, tmp_dir, node)
change_to_agent_session(context)
wait_livebook_app_start(slug)
{:ok, view, _} = live(conn, ~p"/apps")
refute render(view) =~ slug
@ -193,54 +216,4 @@ defmodule LivebookWeb.Integration.AppsLiveTest do
assert render(view) =~ slug
end
end
defp deploy_app(slug, team, org, deployment_group, tmp_dir, node) do
source = """
<!-- livebook:{"app_settings":{"access_type":"public","slug":"#{slug}"},"hub_id":"#{team.id}","deployment_group_id":"#{deployment_group.id}"} -->
# LivebookApp:#{slug}
```elixir
```
"""
{notebook, %{warnings: []}} = Livebook.LiveMarkdown.notebook_from_livemd(source)
files_dir = Livebook.FileSystem.File.local(tmp_dir)
{:ok, %Livebook.Teams.AppDeployment{file: zip_content} = app_deployment} =
Livebook.Teams.AppDeployment.new(notebook, files_dir)
secret_key = Livebook.Teams.derive_key(team.teams_key)
encrypted_content = Livebook.Teams.encrypt(zip_content, secret_key)
app_deployment =
TeamsRPC.upload_app_deployment(
node,
org,
deployment_group,
app_deployment,
encrypted_content,
# broadcast?
true
)
app_deployment_id = to_string(app_deployment.id)
assert_receive {:app_deployment_started, %{id: ^app_deployment_id}}
assert_receive {:app_created, %{pid: pid, slug: ^slug}}, 50000
assert_receive {:app_updated,
%{
slug: ^slug,
sessions: [%{app_status: %{execution: :executed, lifecycle: :active}}]
}},
5000
on_exit(fn ->
if Process.alive?(pid) do
Livebook.App.close(pid)
end
end)
end
end

3
test/livebook_teams/web/hub/edit_live_test.exs
View file

@ -9,7 +9,7 @@ defmodule LivebookWeb.Integration.Hub.EditLiveTest do
setup :teams
@moduletag subscribe_to_hubs_topics: [:crud, :connection, :secrets, :file_systems]
@moduletag subscribe_to_teams_topics: [:clients]
@moduletag subscribe_to_teams_topics: [:clients, :agents]
describe "user" do
@describetag teams_for: :user
@ -286,7 +286,6 @@ defmodule LivebookWeb.Integration.Hub.EditLiveTest do
describe "agent" do
@describetag teams_for: :agent
@describetag subscribe_to_teams_topics: [:clients, :agents]
test "shows an error when creating a secret", %{conn: conn, team: team} do
{:ok, view, _html} = live(conn, ~p"/hub/#{team.id}")

39
test/livebook_teams/zta/livebook_teams_test.exs
View file

@ -10,27 +10,24 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
@moduletag subscribe_to_teams_topics: [:clients, :agents]
describe "authenticate/3" do
setup %{team: team} do
setup %{team: team, test: test} do
Livebook.Apps.subscribe()
start_supervised!(
{Livebook.ZTA.LivebookTeams, name: LivebookWeb.ZTA, identity_key: team.id}
)
start_supervised!({LivebookTeams, name: test, identity_key: team.id})
:ok
end
test "renders HTML with JavaScript redirect", %{conn: conn} do
test "renders HTML with JavaScript redirect", %{conn: conn, test: test} do
conn = init_test_session(conn, %{})
assert {conn, nil} = authenticate(conn, [])
assert {conn, nil} = LivebookTeams.authenticate(test, conn, [])
assert conn.halted
assert html_response(conn, 200) =~ "window.location.href = "
end
test "gets the user information from Livebook Teams", %{conn: conn, node: node} do
test "gets the user information from Livebook Teams", %{conn: conn, node: node, test: test} do
# Step 1: Get redirected to Livebook Teams
conn = init_test_session(conn, %{})
{conn, nil} = authenticate(conn, [])
{conn, nil} = LivebookTeams.authenticate(test, conn, [])
[_, location] = Regex.run(~r/URL\("(.*?)"\)/, html_response(conn, 200))
uri = URI.parse(location)
@ -45,7 +42,7 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
|> init_test_session(Plug.Conn.get_session(conn))
assert {conn, %{id: _id, name: _, email: _, payload: %{}} = metadata} =
authenticate(conn, [])
LivebookTeams.authenticate(test, conn, [])
assert redirected_to(conn, 302) == "/"
@ -54,10 +51,10 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
build_conn(:get, "/")
|> init_test_session(Plug.Conn.get_session(conn))
assert {%{halted: false}, ^metadata} = authenticate(conn, [])
assert {%{halted: false}, ^metadata} = LivebookTeams.authenticate(test, conn, [])
end
test "shows an error when the user does not belong to the org", %{conn: conn} do
test "shows an error when the user does not belong to the org", %{conn: conn, test: test} do
# Step 1: Emulate a request coming from Teams saying the user does belong to the org
conn = init_test_session(conn, %{})
@ -68,7 +65,7 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
conn = %Plug.Conn{conn | params: params_from_teams}
{conn, nil} = authenticate(conn, [])
{conn, nil} = LivebookTeams.authenticate(test, conn, [])
assert conn.status == 302
# Step 2: follow the redirect keeping the session set in previous request
@ -76,7 +73,7 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
build_conn(:get, redirected_to(conn))
|> init_test_session(get_session(conn))
{conn, nil} = authenticate(conn, [])
{conn, nil} = LivebookTeams.authenticate(test, conn, [])
assert html_response(conn, 403) =~
"Failed to authenticate with Livebook Teams: you do not belong to this org"
@ -86,9 +83,9 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
describe "logout/2" do
setup :livebook_teams_auth
test "revoke access token from Livebook Teams", %{conn: conn} do
test "revoke access token from Livebook Teams", %{conn: conn, test: test} do
# Revoke the token and the metadata will be invalid for future requests
assert %{status: 302} = conn = logout(conn)
assert %{status: 302} = conn = LivebookTeams.logout(test, conn)
[url] = get_resp_header(conn, "location")
assert %{status: 200} = Req.get!(url)
@ -97,17 +94,9 @@ defmodule Livebook.ZTA.LivebookTeamsTest do
build_conn(:get, ~p"/")
|> init_test_session(get_session(conn))
{conn, nil} = authenticate(conn, [])
{conn, nil} = LivebookTeams.authenticate(test, conn, [])
assert conn.halted
assert html_response(conn, 200) =~ "window.location.href = "
end
end
defp authenticate(conn, opts) do
LivebookTeams.authenticate(LivebookWeb.ZTA, conn, opts)
end
defp logout(conn) do
LivebookTeams.logout(LivebookWeb.ZTA, conn)
end
end

77
test/support/app_helpers.ex
View file

@ -1,4 +1,8 @@
defmodule Livebook.AppHelpers do
import ExUnit.{Assertions, Callbacks}
alias Livebook.TeamsRPC
def deploy_notebook_sync(notebook) do
app_spec = Livebook.Apps.NotebookAppSpec.new(notebook)
@ -9,7 +13,7 @@ defmodule Livebook.AppHelpers do
{:deploy_result, ^ref, {:ok, pid}} ->
Process.demonitor(ref, [:flush])
ExUnit.Callbacks.on_exit(fn ->
on_exit(fn ->
if Process.alive?(pid) do
Livebook.App.close(pid)
end
@ -18,4 +22,75 @@ defmodule Livebook.AppHelpers do
pid
end
end
def deploy_app(slug, team, org, deployment_group, tmp_dir, node) do
app_path = Path.join(tmp_dir, "#{slug}.livemd")
source =
stamp_notebook(app_path, """
<!-- livebook:{"app_settings":{"access_type":"public","slug":"#{slug}"},"hub_id":"#{team.id}","deployment_group_id":"#{deployment_group.id}"} -->
# LivebookApp:#{slug}
```elixir
IO.puts("Hi")
```
""")
files_dir = Livebook.FileSystem.File.local(tmp_dir)
{:ok, %Livebook.Teams.AppDeployment{file: zip_content} = app_deployment} =
Livebook.Teams.AppDeployment.new(source, files_dir)
secret_key = Livebook.Teams.derive_key(team.teams_key)
encrypted_content = Livebook.Teams.encrypt(zip_content, secret_key)
app_deployment =
TeamsRPC.upload_app_deployment(
node,
org,
deployment_group,
app_deployment,
encrypted_content,
# broadcast?
true
)
app_deployment_id = to_string(app_deployment.id)
assert_receive {:app_deployment_started, %{id: ^app_deployment_id} = app_deployment}
app_deployment
end
def wait_livebook_app_start(slug) do
assert_receive {:app_created, %{pid: pid, slug: ^slug}}, 5_000
assert_receive {:app_updated,
%{
slug: ^slug,
sessions: [%{app_status: %{execution: :executed, lifecycle: :active}}]
}},
5_000
on_exit(fn ->
if Process.alive?(pid) do
Livebook.App.close(pid)
Process.exit(pid, :kill)
end
end)
pid
end
def stamp_notebook(app_path, %Livebook.Notebook{} = notebook) do
{source, []} = Livebook.LiveMarkdown.notebook_to_livemd(notebook)
File.write!(app_path, source)
source
end
def stamp_notebook(app_path, source) when is_binary(source) do
{notebook, %{warnings: []}} = Livebook.LiveMarkdown.notebook_from_livemd(source)
stamp_notebook(app_path, notebook)
end
end

126
test/support/integration/teams_tests.ex
View file

@ -1,5 +1,5 @@
defmodule Livebook.TeamsIntegrationHelper do
alias Livebook.{Factory, Hubs, Teams, TeamsRPC}
alias Livebook.{Factory, Hubs, Teams, TeamsRPC, ZTA}
import ExUnit.Assertions
import Phoenix.ConnTest
@ -16,29 +16,19 @@ defmodule Livebook.TeamsIntegrationHelper do
end
end
def livebook_teams_auth(%{conn: conn, node: node, team: team, teams_for: :agent} = context) do
ExUnit.Callbacks.start_supervised!(
{Livebook.ZTA.LivebookTeams, name: LivebookWeb.ZTA, identity_key: team.id}
)
def livebook_teams_auth(%{conn: conn, 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(conn, node, team)
Map.merge(context, %{conn: conn, code: code})
end
@doc false
def create_user_hub(node) do
context = new_user_hub(node)
id = context.team.id
Hubs.save_hub(context.team)
pid = Hubs.TeamClient.get_pid(id)
assert Process.alive?(pid)
assert Hubs.hub_exists?(id)
assert_receive {:hub_connected, ^id}, 3_000
assert_receive {:client_connected, ^id}, 3_000
ExUnit.Callbacks.on_exit(fn -> Hubs.delete_hub(id) end)
context
wait_until_client_start(context)
end
@doc false
@ -56,6 +46,7 @@ defmodule Livebook.TeamsIntegrationHelper do
Factory.build(:team,
id: "team-#{org.name}",
hub_name: org.name,
hub_emoji: "💡",
user_id: user.id,
org_id: org.id,
org_key_id: org_key.id,
@ -70,6 +61,7 @@ defmodule Livebook.TeamsIntegrationHelper do
user: user,
org_key: org_key,
org_key_pair: org_key_pair,
session_token: token,
team: team
}
end
@ -77,25 +69,9 @@ defmodule Livebook.TeamsIntegrationHelper do
@doc false
def create_agent_hub(node, opts \\ []) do
context = new_agent_hub(node, opts)
id = context.team.id
Hubs.save_hub(context.team)
deployment_group_id = to_string(context.deployment_group.id)
org_id = to_string(context.org.id)
pid = Hubs.TeamClient.get_pid(id)
assert Process.alive?(pid)
assert Hubs.hub_exists?(id)
assert_receive {:hub_connected, ^id}, 3_000
assert_receive {:client_connected, ^id}, 3_000
assert_receive {:agent_joined,
%{hub_id: ^id, deployment_group_id: ^deployment_group_id, org_id: ^org_id} =
agent},
3_000
ExUnit.Callbacks.on_exit(fn -> Hubs.delete_hub(id) end)
Map.put_new(%{context | team: Hubs.fetch_hub!(id)}, :agent, agent)
wait_until_agent_start(context)
end
@doc false
@ -124,6 +100,7 @@ defmodule Livebook.TeamsIntegrationHelper do
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,
@ -143,10 +120,13 @@ defmodule Livebook.TeamsIntegrationHelper do
end
@doc false
def authenticate_user_on_teams(conn, node, team) do
def authenticate_user_on_teams(name, conn, node, team) do
# Create a fresh connection to avoid session contamination
fresh_conn = Phoenix.ConnTest.build_conn()
response =
conn
|> LivebookWeb.ConnCase.with_authorization(team.id)
fresh_conn
|> LivebookWeb.ConnCase.with_authorization(team.id, name)
|> get("/")
|> html_response(200)
@ -157,20 +137,84 @@ defmodule Livebook.TeamsIntegrationHelper do
%{code: code} = Livebook.TeamsRPC.allow_auth_request(node, token)
session =
conn
|> LivebookWeb.ConnCase.with_authorization(team.id)
fresh_conn
|> LivebookWeb.ConnCase.with_authorization(team.id, name)
|> get("/", %{teams_identity: "", code: code})
|> Plug.Conn.get_session()
conn = Plug.Test.init_test_session(conn, session)
authenticated_conn = get(conn, "/")
assigns = Map.take(authenticated_conn.assigns, [:current_user])
# 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])
{%Plug.Conn{conn | assigns: Map.merge(conn.assigns, assigns)}, code}
{%Plug.Conn{authenticated_conn | assigns: Map.merge(authenticated_conn.assigns, assigns)},
code}
end
@doc false
def change_to_agent_session(%{node: node, teams_for: :user} = context) do
pid = Hubs.TeamClient.get_pid(context.team.id)
Hubs.TeamClient.stop(context.team.id)
refute Process.alive?(pid)
agent_key = context[:agent_key] || TeamsRPC.create_agent_key(node, org: context.org)
deployment_group =
context[:deployment_group] ||
TeamsRPC.create_deployment_group(node, mode: :online, org: context.org)
team = %{context.team | user_id: nil, session_token: agent_key.key}
Hubs.save_hub(team)
%{context | teams_for: :agent}
|> Map.put_new(:agent_key, agent_key)
|> Map.put_new(:deployment_group, deployment_group)
|> wait_until_agent_start()
end
def change_to_user_session(%{node: node, org: org, teams_for: :agent} = context) do
pid = Hubs.TeamClient.get_pid(context.team.id)
Hubs.TeamClient.stop(context.team.id)
refute Process.alive?(pid)
user = context[:user] || TeamsRPC.create_user(node)
session_token = context[:session_token] || TeamsRPC.associate_user_with_org(node, user, org)
team = %{context.team | user_id: user.id, session_token: session_token}
Hubs.save_hub(team)
wait_until_client_start(%{context | team: team, teams_for: :user})
end
# Private
defp wait_until_client_start(context) do
id = context.team.id
pid = Hubs.TeamClient.get_pid(id)
assert Process.alive?(pid)
assert Hubs.hub_exists?(id)
assert_receive {:hub_connected, ^id}, 3_000
assert_receive {:client_connected, ^id}, 3_000
context
end
defp wait_until_agent_start(context) do
context = wait_until_client_start(context)
id = context.team.id
deployment_group_id = to_string(context.deployment_group.id)
org_id = to_string(context.org.id)
assert_receive {:agent_joined,
%{hub_id: ^id, deployment_group_id: ^deployment_group_id, org_id: ^org_id} =
agent},
3_000
Map.put_new(context, :agent, agent)
end
defp generate_key_hash(teams_key \\ Teams.Org.teams_key()) do
{teams_key, Teams.Org.key_hash(%Teams.Org{teams_key: teams_key})}
end