Make feature toggle work with teams schema and adjust affected tests (#…

…4825)

lib/plausible/billing/feature.ex

@@ -48,7 +48,8 @@ defmodule Plausible.Billing.Feature do
   `{:error, :upgrade_required}` when toggling a feature the site owner does not
   have access to.
   """
-  @callback toggle(Plausible.Site.t(), Keyword.t()) :: :ok | {:error, :upgrade_required}
+  @callback toggle(Plausible.Site.t(), Plausible.Auth.User.t(), Keyword.t()) ::
+              :ok | {:error, :upgrade_required}
 
   @doc """
   Checks whether a feature is enabled or not. Returns false when the feature is
@@ -130,31 +131,20 @@ defmodule Plausible.Billing.Feature do
 
       @impl true
       def check_availability(%Plausible.Auth.User{} = user) do
-        cond do
-          free?() -> :ok
-          __MODULE__ in Quota.Limits.allowed_features_for(user) -> :ok
-          true -> {:error, :upgrade_required}
-        end
-      end
-
-      def check_availability(team_or_nil) do
-        cond do
-          free?() -> :ok
-          __MODULE__ in Plausible.Teams.Billing.allowed_features_for(team_or_nil) -> :ok
-          true -> {:error, :upgrade_required}
-        end
+        Plausible.Teams.Adapter.Read.Billing.check_feature_availability(__MODULE__, user)
       end
 
       @impl true
-      def toggle(%Plausible.Site{} = site, opts \\ []) do
-        if toggle_field(), do: do_toggle(site, opts), else: :ok
+      def toggle(%Plausible.Site{} = site, %Plausible.Auth.User{} = user, opts \\ []) do
+        if toggle_field(), do: do_toggle(site, user, opts), else: :ok
       end
 
-      defp do_toggle(%Plausible.Site{} = site, opts) do
-        site = Plausible.Repo.preload(site, :owner)
+      defp do_toggle(%Plausible.Site{} = site, user, opts) do
+        owner = Plausible.Teams.Adapter.Read.Ownership.get_owner(site, user)
+
         override = Keyword.get(opts, :override)
         toggle = if is_boolean(override), do: override, else: !Map.fetch!(site, toggle_field())
-        availability = if toggle, do: check_availability(site.owner), else: :ok
+        availability = if toggle, do: check_availability(owner), else: :ok
 
         case availability do
           :ok ->
@@ -210,38 +200,17 @@ defmodule Plausible.Billing.Feature.StatsAPI do
     name: :stats_api,
     display_name: "Stats API"
 
-  if Plausible.ee?() do
-    @impl true
-    @doc """
-    Checks whether the user has access to Stats API or not.
-
-    Before the business tier, users who had not yet started their trial had
-    access to Stats API. With the business tier work, access is blocked and they
-    must either start their trial or subscribe to a plan. This is common when a
-    site owner invites a new user. In such cases, using the owner's API key is
-    recommended.
-    """
-    def check_availability(%Plausible.Auth.User{} = user) do
-      user = Plausible.Users.with_subscription(user)
-      unlimited_trial? = is_nil(user.trial_expiry_date)
-      subscription? = Plausible.Billing.Subscriptions.active?(user.subscription)
-
-      pre_business_tier_account? =
-        NaiveDateTime.before?(user.inserted_at, Plausible.Billing.Plans.business_tier_launch())
-
-      cond do
-        !subscription? && unlimited_trial? && pre_business_tier_account? ->
-          :ok
-
-        !subscription? && unlimited_trial? && !pre_business_tier_account? ->
-          {:error, :upgrade_required}
-
-        true ->
-          super(user)
-      end
-    end
-  else
-    @impl true
-    def check_availability(_user), do: :ok
+  @impl true
+  @doc """
+  Checks whether the user has access to Stats API or not.
+
+  Before the business tier, users who had not yet started their trial had
+  access to Stats API. With the business tier work, access is blocked and they
+  must either start their trial or subscribe to a plan. This is common when a
+  site owner invites a new user. In such cases, using the owner's API key is
+  recommended.
+  """
+  def check_availability(%Plausible.Auth.User{} = user) do
+    Plausible.Teams.Adapter.Read.Billing.check_feature_availability_for_stats_api(user)
   end
 end

lib/plausible/teams/adapter.ex

@@ -11,14 +11,6 @@ defmodule Plausible.Teams.Adapter do
     end
   end
 
-  def team_or_user(user) do
-    switch(
-      user,
-      team_fn: &Function.identity/1,
-      user_fn: &Function.identity/1
-    )
-  end
-
   def switch(user, opts \\ []) do
     team_fn = Keyword.fetch!(opts, :team_fn)
     user_fn = Keyword.fetch!(opts, :user_fn)

lib/plausible/teams/adapter/read/billing.ex

@@ -1,6 +1,6 @@
 defmodule Plausible.Teams.Adapter.Read.Billing do
   @moduledoc """
-  Transition adapter for new schema reads 
+  Transition adapter for new schema reads
   """
   use Plausible.Teams.Adapter
 
@@ -34,4 +34,65 @@ defmodule Plausible.Teams.Adapter.Read.Billing do
       user_fn: &Plausible.Billing.Quota.Usage.site_usage/1
     )
   end
+
+  use Plausible
+
+  on_ee do
+    def check_feature_availability_for_stats_api(user) do
+      {unlimited_trial?, subscription?} =
+        switch(user,
+          team_fn: fn team ->
+            team = Plausible.Teams.with_subscription(team)
+            unlimited_trial? = is_nil(team) or is_nil(team.trial_expiry_date)
+
+            subscription? =
+              not is_nil(team) and Plausible.Billing.Subscriptions.active?(team.subscription)
+
+            {unlimited_trial?, subscription?}
+          end,
+          user_fn: fn user ->
+            user = Plausible.Users.with_subscription(user)
+            unlimited_trial? = is_nil(user.trial_expiry_date)
+            subscription? = Plausible.Billing.Subscriptions.active?(user.subscription)
+
+            {unlimited_trial?, subscription?}
+          end
+        )
+
+      pre_business_tier_account? =
+        NaiveDateTime.before?(user.inserted_at, Plausible.Billing.Plans.business_tier_launch())
+
+      cond do
+        !subscription? && unlimited_trial? && pre_business_tier_account? ->
+          :ok
+
+        !subscription? && unlimited_trial? && !pre_business_tier_account? ->
+          {:error, :upgrade_required}
+
+        true ->
+          check_feature_availability(Plausible.Billing.Feature.StatsAPI, user)
+      end
+    end
+  else
+    def check_feature_availability_for_stats_api(_user), do: :ok
+  end
+
+  def check_feature_availability(feature, user) do
+    switch(user,
+      team_fn: fn team_or_nil ->
+        cond do
+          feature.free?() -> :ok
+          feature in Teams.Billing.allowed_features_for(team_or_nil) -> :ok
+          true -> {:error, :upgrade_required}
+        end
+      end,
+      user_fn: fn user ->
+        cond do
+          feature.free?() -> :ok
+          feature in Plausible.Billing.Quota.Limits.allowed_features_for(user) -> :ok
+          true -> {:error, :upgrade_required}
+        end
+      end
+    )
+  end
 end

lib/plausible/teams/adapter/read/ownership.ex

@@ -1,13 +1,27 @@
 defmodule Plausible.Teams.Adapter.Read.Ownership do
   @moduledoc """
-  Transition adapter for new schema reads 
+  Transition adapter for new schema reads
   """
   use Plausible
   use Plausible.Teams.Adapter
   alias Plausible.Site
   alias Plausible.Auth
   alias Plausible.Site.Memberships.Invitations
 
+  def get_owner(site, user) do
+    switch(user,
+      team_fn: fn team ->
+        case Teams.Sites.get_owner(team) do
+          {:ok, user} -> user
+          _ -> nil
+        end
+      end,
+      user_fn: fn _ ->
+        Plausible.Repo.preload(site, :owner).owner
+      end
+    )
+  end
+
   def ensure_can_take_ownership(site, user) do
     switch(
       user,
@@ -39,11 +53,9 @@ defmodule Plausible.Teams.Adapter.Read.Ownership do
 
   on_ee do
     def check_feature_access(site, new_owner) do
-      team_or_user = team_or_user(new_owner)
-
       missing_features =
         Plausible.Billing.Quota.Usage.features_usage(nil, [site.id])
-        |> Enum.filter(&(&1.check_availability(team_or_user) != :ok))
+        |> Enum.filter(&(&1.check_availability(new_owner) != :ok))
 
       if missing_features == [] do
         :ok

lib/plausible_web/controllers/api/internal_controller.ex

@@ -25,11 +25,11 @@ defmodule PlausibleWeb.Api.InternalController do
     "conversions" => Plausible.Billing.Feature.Goals
   }
   def disable_feature(conn, %{"domain" => domain, "feature" => feature}) do
-    with %User{id: user_id} <- conn.assigns[:current_user],
+    with %User{id: user_id} = user <- conn.assigns[:current_user],
          site <- Sites.get_by_domain(domain),
          true <- Sites.has_admin_access?(user_id, site) || Auth.is_super_admin?(user_id),
          {:ok, mod} <- Map.fetch(@features, feature),
-         {:ok, _site} <- mod.toggle(site, override: false) do
+         {:ok, _site} <- mod.toggle(site, user, override: false) do
       json(conn, "ok")
     else
       {:error, :upgrade_required} ->

lib/plausible_web/controllers/site_controller.ex

@@ -81,7 +81,7 @@ defmodule PlausibleWeb.SiteController do
     feature_mod =
       Enum.find(Plausible.Billing.Feature.list(), &(&1.toggle_field() == toggle_field))
 
-    case feature_mod.toggle(site, override: value == "true") do
+    case feature_mod.toggle(site, conn.assigns.current_user, override: value == "true") do
       {:ok, updated_site} ->
         message =
           if Map.fetch!(updated_site, toggle_field) do

test/plausible/auth/auth_test.exs

@@ -1,5 +1,6 @@
 defmodule Plausible.AuthTest do
   use Plausible.DataCase, async: true
+  use Plausible.Teams.Test
   alias Plausible.Auth
 
   describe "user_completed_setup?" do
@@ -52,14 +53,14 @@ defmodule Plausible.AuthTest do
 
   describe "create_api_key/3" do
     test "creates a new api key" do
-      user = insert(:user)
+      user = new_user()
       key = Ecto.UUID.generate()
       assert {:ok, %Auth.ApiKey{}} = Auth.create_api_key(user, "my new key", key)
     end
 
     @tag :ee_only
     test "defaults to 600 requests per hour limit in EE" do
-      user = insert(:user)
+      user = new_user()
 
       {:ok, %Auth.ApiKey{hourly_request_limit: hourly_request_limit}} =
         Auth.create_api_key(user, "my new EE key", Ecto.UUID.generate())
@@ -78,8 +79,8 @@ defmodule Plausible.AuthTest do
     end
 
     test "errors when key already exists" do
-      u1 = insert(:user)
-      u2 = insert(:user)
+      u1 = new_user()
+      u2 = new_user()
       key = Ecto.UUID.generate()
       assert {:ok, %Auth.ApiKey{}} = Auth.create_api_key(u1, "my new key", key)
       assert {:error, changeset} = Auth.create_api_key(u2, "my other key", key)
@@ -100,16 +101,16 @@ defmodule Plausible.AuthTest do
 
   describe "delete_api_key/2" do
     test "deletes the record" do
-      user = insert(:user)
+      user = new_user()
       assert {:ok, api_key} = Auth.create_api_key(user, "my new key", Ecto.UUID.generate())
       assert :ok = Auth.delete_api_key(user, api_key.id)
       refute Plausible.Repo.reload(api_key)
     end
 
     test "returns error when api key does not exist or does not belong to user" do
-      me = insert(:user)
+      me = new_user()
 
-      other_user = insert(:user)
+      other_user = new_user()
       {:ok, other_api_key} = Auth.create_api_key(other_user, "my new key", Ecto.UUID.generate())
 
       assert {:error, :not_found} = Auth.delete_api_key(me, other_api_key.id)