Fix tests

Author: aleDsz

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

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)

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 """

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]

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 =

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)

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",