use in timelines benchmark new user generator

3 files changed, 14 insertions, 19 deletions

diff --git a/benchmarks/load_testing/activities.ex b/benchmarks/load_testing/activities.ex
index 24c6b5531..23ee2b987 100644
--- a/benchmarks/load_testing/activities.ex
+++ b/benchmarks/load_testing/activities.ex
@@ -81,7 +81,7 @@ defmodule Pleroma.LoadTesting.Activities do
         )
       end)
 
-    IO.puts("Generating iterations of activities take #{to_sec(time)} sec.\n")
+    IO.puts("Generating iterations of activities took #{to_sec(time)} sec.\n")
     :ok
   end
 
diff --git a/benchmarks/load_testing/users.ex b/benchmarks/load_testing/users.ex
index b73ac8651..1a8c6e22f 100644
--- a/benchmarks/load_testing/users.ex
+++ b/benchmarks/load_testing/users.ex
@@ -33,7 +33,7 @@ defmodule Pleroma.LoadTesting.Users do
   def generate_users(max) do
     IO.puts("Starting generating #{max} users...")
 
-    {time, _} =
+    {time, users} =
       :timer.tc(fn ->
         Task.async_stream(
           1..max,
@@ -41,10 +41,11 @@ defmodule Pleroma.LoadTesting.Users do
           max_concurrency: @max_concurrency,
           timeout: 30_000
         )
-        |> Stream.run()
+        |> Enum.to_list()
       end)
 
-    IO.puts("Generating users take #{to_sec(time)} sec.\n")
+    IO.puts("Generating users took #{to_sec(time)} sec.\n")
+    users
   end
 
   defp generate_user(i) do
@@ -104,7 +105,7 @@ defmodule Pleroma.LoadTesting.Users do
         |> run_stream(main_user)
       end)
 
-    IO.puts("Making friends take #{to_sec(time)} sec.\n")
+    IO.puts("Making friends took #{to_sec(time)} sec.\n")
   end
 
   def make_friends(%User{} = main_user, %User{} = user) do
diff --git a/benchmarks/mix/tasks/pleroma/benchmarks/timelines.ex b/benchmarks/mix/tasks/pleroma/benchmarks/timelines.ex
index dc6f3d3fc..9b7ac6111 100644
--- a/benchmarks/mix/tasks/pleroma/benchmarks/timelines.ex
+++ b/benchmarks/mix/tasks/pleroma/benchmarks/timelines.ex
@@ -1,9 +1,10 @@
 defmodule Mix.Tasks.Pleroma.Benchmarks.Timelines do
   use Mix.Task
-  alias Pleroma.Repo
-  alias Pleroma.LoadTesting.Generator
+
+  import Pleroma.LoadTesting.Helper, only: [clean_tables: 0]
 
   alias Pleroma.Web.CommonAPI
+  alias Plug.Conn
 
   def run(_args) do
     Mix.Pleroma.start_pleroma()
@@ -11,7 +12,7 @@ defmodule Mix.Tasks.Pleroma.Benchmarks.Timelines do
     # Cleaning tables
     clean_tables()
 
-    [{:ok, user} | users] = Generator.generate_users(users_max: 1000)
+    [{:ok, user} | users] = Pleroma.LoadTesting.Users.generate_users(1000)
 
     # Let the user make 100 posts
 
@@ -38,8 +39,8 @@ defmodule Mix.Tasks.Pleroma.Benchmarks.Timelines do
         "user timeline, no followers" => fn reading_user ->
           conn =
             Phoenix.ConnTest.build_conn()
-            |> Plug.Conn.assign(:user, reading_user)
-            |> Plug.Conn.assign(:skip_link_headers, true)
+            |> Conn.assign(:user, reading_user)
+            |> Conn.assign(:skip_link_headers, true)
 
           Pleroma.Web.MastodonAPI.AccountController.statuses(conn, %{"id" => user.id})
         end
@@ -56,8 +57,8 @@ defmodule Mix.Tasks.Pleroma.Benchmarks.Timelines do
         "user timeline, all following" => fn reading_user ->
           conn =
             Phoenix.ConnTest.build_conn()
-            |> Plug.Conn.assign(:user, reading_user)
-            |> Plug.Conn.assign(:skip_link_headers, true)
+            |> Conn.assign(:user, reading_user)
+            |> Conn.assign(:skip_link_headers, true)
 
           Pleroma.Web.MastodonAPI.AccountController.statuses(conn, %{"id" => user.id})
         end
@@ -66,11 +67,4 @@ defmodule Mix.Tasks.Pleroma.Benchmarks.Timelines do
       time: 60
     )
   end
-
-  defp clean_tables do
-    IO.puts("Deleting old data...\n")
-    Ecto.Adapters.SQL.query!(Repo, "TRUNCATE users CASCADE;")
-    Ecto.Adapters.SQL.query!(Repo, "TRUNCATE activities CASCADE;")
-    Ecto.Adapters.SQL.query!(Repo, "TRUNCATE objects CASCADE;")
-  end
 end