From 5b0008130f355908c110d33e7f1db7da1b1da94d Mon Sep 17 00:00:00 2001
From: TzeYiing <ty@tzeyiing.com>
Date: Tue, 14 Nov 2023 17:11:43 +0800
Subject: [PATCH] chore: add specs, minor refactoring of Logflare.Cluster.Utils

---
 lib/logflare/cluster/utils.ex  | 16 ++++++++++------
 test/logflare/cluster_test.exs | 28 ++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 6 deletions(-)
 create mode 100644 test/logflare/cluster_test.exs

diff --git a/lib/logflare/cluster/utils.ex b/lib/logflare/cluster/utils.ex
index 11aaca0e0..bd1d4dbaf 100644
--- a/lib/logflare/cluster/utils.ex
+++ b/lib/logflare/cluster/utils.ex
@@ -2,27 +2,31 @@ defmodule Logflare.Cluster.Utils do
   @moduledoc false
   require Logger
 
-  defp env_min_cluster_size, do: Application.get_env(:logflare, __MODULE__)[:min_cluster_size]
-
+  @spec node_list_all() :: [Node.t()]
   def node_list_all() do
     [Node.self() | Node.list()]
   end
 
+  @spec cluster_size() :: non_neg_integer()
   def cluster_size() do
-    lib_cluster_size = node_list_all() |> Enum.count()
+    lib_cluster_size = actual_cluster_size()
+    min_size = env_min_cluster_size()
 
-    if lib_cluster_size >= env_min_cluster_size() do
+    if lib_cluster_size >= min_size do
       lib_cluster_size
     else
-      Logger.warning("Cluster size is #{lib_cluster_size} but expected #{env_min_cluster_size()}",
+      Logger.warning("Cluster size is #{lib_cluster_size} but expected #{min_size}",
         cluster_size: lib_cluster_size
       )
 
-      env_min_cluster_size()
+      min_size
     end
   end
 
+  @spec actual_cluster_size() :: non_neg_integer()
   def actual_cluster_size() do
     Enum.count(node_list_all())
   end
+
+  defp env_min_cluster_size, do: Application.get_env(:logflare, __MODULE__)[:min_cluster_size]
 end
diff --git a/test/logflare/cluster_test.exs b/test/logflare/cluster_test.exs
new file mode 100644
index 000000000..e18747454
--- /dev/null
+++ b/test/logflare/cluster_test.exs
@@ -0,0 +1,28 @@
+defmodule Logflare.ClusterTest do
+  use ExUnit.Case, async: false
+  alias Logflare.Cluster
+
+  test "Cluster.Utils.node_list_all/0" do
+    assert [_|_] = Cluster.Utils.node_list_all()
+  end
+
+
+  test "cluster_size/0, actual_cluster_size/0" do
+    assert Cluster.Utils.cluster_size() == Cluster.Utils.actual_cluster_size()
+  end
+
+  describe "Cluster.Utils" do
+    setup do
+      config = Application.get_env(:logflare, Cluster.Utils)
+      Application.put_env(:logflare, Cluster.Utils, [min_cluster_size: 12])
+      on_exit(fn ->
+        Application.put_env(:logflare, Cluster.Utils, config)
+      end)
+    end
+    test "cluster_size/0, actual_cluster_size/0" do
+      assert Cluster.Utils.cluster_size() == 12
+      assert Cluster.Utils.actual_cluster_size() == 1
+    end
+
+  end
+end