From a8e085037e38e8c555169c556b42d7efa5416688 Mon Sep 17 00:00:00 2001 From: Alexandre de Souza Date: Tue, 15 Jul 2025 17:00:06 -0300 Subject: [PATCH] Fix tests --- lib/livebook/hubs.ex | 34 +++-- lib/livebook/hubs/team.ex | 5 +- lib/livebook/hubs/team_client.ex | 20 +++ test/livebook_teams/apps_test.exs | 2 +- test/livebook_teams/hubs/team_client_test.exs | 55 ++++++-- test/livebook_teams/teams_test.exs | 13 +- test/livebook_teams/web/admin_live_test.exs | 2 +- .../web/app_session_live_test.exs | 92 +++++-------- test/livebook_teams/web/apps_live_test.exs | 85 ++++-------- .../livebook_teams/web/hub/edit_live_test.exs | 3 +- .../zta/livebook_teams_test.exs | 39 ++---- test/support/app_helpers.ex | 77 ++++++++++- test/support/integration/teams_tests.ex | 126 ++++++++++++------ 13 files changed, 338 insertions(+), 215 deletions(-) diff --git a/lib/livebook/hubs.ex b/lib/livebook/hubs.ex index 1c92aae5f..ef0829103 100644 --- a/lib/livebook/hubs.ex +++ b/lib/livebook/hubs.ex @@ -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 diff --git a/lib/livebook/hubs/team.ex b/lib/livebook/hubs/team.ex index 561963f27..7832557af 100644 --- a/lib/livebook/hubs/team.ex +++ b/lib/livebook/hubs/team.ex @@ -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) diff --git a/lib/livebook/hubs/team_client.ex b/lib/livebook/hubs/team_client.ex index 1ea11041d..a578b3336 100644 --- a/lib/livebook/hubs/team_client.ex +++ b/lib/livebook/hubs/team_client.ex @@ -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 """ diff --git a/test/livebook_teams/apps_test.exs b/test/livebook_teams/apps_test.exs index 292e55cc1..1193c9b75 100644 --- a/test/livebook_teams/apps_test.exs +++ b/test/livebook_teams/apps_test.exs @@ -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] diff --git a/test/livebook_teams/hubs/team_client_test.exs b/test/livebook_teams/hubs/team_client_test.exs index f3b21249f..197ebf164 100644 --- a/test/livebook_teams/hubs/team_client_test.exs +++ b/test/livebook_teams/hubs/team_client_test.exs @@ -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 = diff --git a/test/livebook_teams/teams_test.exs b/test/livebook_teams/teams_test.exs index c877bb219..6261a695a 100644 --- a/test/livebook_teams/teams_test.exs +++ b/test/livebook_teams/teams_test.exs @@ -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) diff --git a/test/livebook_teams/web/admin_live_test.exs b/test/livebook_teams/web/admin_live_test.exs index 26705457a..1aecdcfed 100644 --- a/test/livebook_teams/web/admin_live_test.exs +++ b/test/livebook_teams/web/admin_live_test.exs @@ -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", diff --git a/test/livebook_teams/web/app_session_live_test.exs b/test/livebook_teams/web/app_session_live_test.exs index b6d32733d..edd97ce57 100644 --- a/test/livebook_teams/web/app_session_live_test.exs +++ b/test/livebook_teams/web/app_session_live_test.exs @@ -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 = """ - - - # 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 diff --git a/test/livebook_teams/web/apps_live_test.exs b/test/livebook_teams/web/apps_live_test.exs index 4ddc2870d..746c1b391 100644 --- a/test/livebook_teams/web/apps_live_test.exs +++ b/test/livebook_teams/web/apps_live_test.exs @@ -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 = """ - - - # 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 diff --git a/test/livebook_teams/web/hub/edit_live_test.exs b/test/livebook_teams/web/hub/edit_live_test.exs index 069656dad..6f8444407 100644 --- a/test/livebook_teams/web/hub/edit_live_test.exs +++ b/test/livebook_teams/web/hub/edit_live_test.exs @@ -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}") diff --git a/test/livebook_teams/zta/livebook_teams_test.exs b/test/livebook_teams/zta/livebook_teams_test.exs index 44e092ce1..e1efc2588 100644 --- a/test/livebook_teams/zta/livebook_teams_test.exs +++ b/test/livebook_teams/zta/livebook_teams_test.exs @@ -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 diff --git a/test/support/app_helpers.ex b/test/support/app_helpers.ex index 56ba966c6..f32551748 100644 --- a/test/support/app_helpers.ex +++ b/test/support/app_helpers.ex @@ -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, """ + + + # 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 diff --git a/test/support/integration/teams_tests.ex b/test/support/integration/teams_tests.ex index 3e8b28d50..64cba985e 100644 --- a/test/support/integration/teams_tests.ex +++ b/test/support/integration/teams_tests.ex @@ -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