Add more tests for KafkaUtils

lib/kafkaex_lag_exporter/behaviours/kafka_utils_behaviour.ex

@@ -3,16 +3,6 @@ defmodule KafkaexLagExporter.KafkaUtils.Behaviour do
 
   @callback connection(atom) :: {list({charlist, non_neg_integer}), Keyword.t()}
 
-  @callback resolve_offsets(binary, :earliest | :latest, atom) ::
-              list({non_neg_integer, integer})
-
-  @callback fetch_committed_offsets(binary, binary, atom) ::
-              list({non_neg_integer, non_neg_integer})
-
-  @callback lag(binary, binary, atom) :: list({non_neg_integer, integer})
-
-  @callback lag_total(binary, binary, atom) :: non_neg_integer
-
   @callback get_consumer_group_names({host :: atom, port :: non_neg_integer}) :: list(binary)
 
   @callback get_consumer_group_info(
@@ -21,23 +11,20 @@ defmodule KafkaexLagExporter.KafkaUtils.Behaviour do
               list(binary)
             ) :: list({consumer_group :: binary, topics :: list(binary)})
 
+  @callback lag(binary, binary, atom) :: list({non_neg_integer, integer})
+
   def connection(client), do: impl().connection(client)
 
-  def resolve_offsets(topic, type, client), do: impl().resolve_offsets(topic, type, client)
-
-  def fetch_committed_offsets(topic, consumer_group, client),
-    do: impl().fetch_committed_offsets(topic, consumer_group, client)
-
-  def lag(topic, consumer_group, client), do: impl().lag(topic, consumer_group, client)
-
-  def lag_total(topic, consumer_group, client),
-    do: impl().lag_total(topic, consumer_group, client)
-
-  def get_consumer_group_names({host, port}), do: impl().get_consumer_group_names({host, port})
+  def get_consumer_group_names(endpoint), do: impl().get_consumer_group_names(endpoint)
 
   def get_consumer_group_info(endpoint, list, consumer_group_names),
     do: impl().get_consumer_group_info(endpoint, list, consumer_group_names)
 
+  def lag(topic, consumer_group, client), do: impl().lag(topic, consumer_group, client)
+
+  def fetch_committed_offsets(topic, consumer_group, client),
+    do: impl().fetch_committed_offsets(topic, consumer_group, client)
+
   defp impl,
     do: Application.get_env(:kafkaex_lag_exporter, :kafka_utils, KafkaexLagExporter.KafkaUtils)
 end

lib/kafkaex_lag_exporter/kafka_utils.ex

@@ -32,54 +32,8 @@ defmodule KafkaexLagExporter.KafkaUtils do
   end
 
   @impl true
-  def resolve_offsets(topic, type, client) do
-    {endpoints, sock_opts} = connection(client)
-
-    {:ok, partitions_count} = KafkaWrapper.get_partitions_count(client, topic)
-
-    for i <- Range.new(0, partitions_count - 1),
-        {:ok, offset} =
-          KafkaWrapper.resolve_offset(endpoints, topic, i, type, sock_opts) do
-      {i, offset}
-    end
-  end
-
-  @impl true
-  def fetch_committed_offsets(_topic, consumer_group, client) do
-    {endpoints, sock_opts} = connection(client)
-
-    {:ok, response} = KafkaWrapper.fetch_committed_offsets(endpoints, sock_opts, consumer_group)
-
-    for r <- response,
-        pr <- r[:partitions],
-        do: {pr[:partition_index], pr[:committed_offset]}
-  end
-
-  @impl true
-  def lag(topic, consumer_group, client) do
-    offsets =
-      resolve_offsets(topic, :latest, client)
-      |> Enum.sort_by(fn {key, _value} -> key end)
-
-    committed_offsets =
-      fetch_committed_offsets(topic, consumer_group, client)
-      |> Enum.sort_by(fn {key, _value} -> key end)
-
-    for {{part, current}, {_part2, committed}} <- Enum.zip(offsets, committed_offsets) do
-      {part, current - committed}
-    end
-  end
-
-  @impl true
-  def lag_total(topic, consumer_group, client) do
-    for {_part, recs} <- lag(topic, consumer_group, client), reduce: 0 do
-      acc -> acc + recs
-    end
-  end
-
-  @impl true
-  def get_consumer_group_names({host, port}) do
-    [{_, groups} | _] = KafkaWrapper.list_all_groups([{host, port}], [])
+  def get_consumer_group_names(endpoint) do
+    [{_, groups} | _] = KafkaWrapper.list_all_groups([endpoint], [])
 
     groups
     |> Enum.filter(fn {_, _, protocol} -> protocol === "consumer" end)
@@ -93,12 +47,52 @@ defmodule KafkaexLagExporter.KafkaUtils do
     group_descriptions
     |> Enum.flat_map(fn %{group_id: consumer_group, members: members} ->
       get_member_info(members)
-      |> Enum.map(fn {topic_names, consumer_id, member_host} ->
-        {consumer_group, topic_names, consumer_id, member_host}
+      |> Enum.map(fn {topics, consumer_id, member_host} ->
+        {consumer_group, topics, consumer_id, member_host}
       end)
     end)
   end
 
+  @impl true
+  def lag(topic, consumer_group, client) do
+    offsets =
+      resolve_offsets(topic, client)
+      |> Enum.sort_by(fn {key, _value} -> key end)
+
+    committed_offsets =
+      fetch_committed_offsets(topic, consumer_group, client)
+      |> Enum.sort_by(fn {key, _value} -> key end)
+
+    for {{part, current}, {_part2, committed}} <- Enum.zip(offsets, committed_offsets) do
+      {part, current - committed}
+    end
+  end
+
+  @spec resolve_offsets(binary, atom) :: list({non_neg_integer, integer})
+  defp resolve_offsets(topic, client) do
+    {endpoints, sock_opts} = connection(client)
+
+    {:ok, partitions_count} = KafkaWrapper.get_partitions_count(client, topic)
+
+    for i <- Range.new(0, partitions_count - 1),
+        {:ok, offset} =
+          KafkaWrapper.resolve_offset(endpoints, topic, i, :latest, sock_opts) do
+      {i, offset}
+    end
+  end
+
+  @spec fetch_committed_offsets(binary, binary, atom) ::
+          list({non_neg_integer, non_neg_integer})
+  defp fetch_committed_offsets(_topic, consumer_group, client) do
+    {endpoints, sock_opts} = connection(client)
+
+    {:ok, response} = KafkaWrapper.fetch_committed_offsets(endpoints, sock_opts, consumer_group)
+
+    for r <- response,
+        pr <- r[:partitions],
+        do: {pr[:partition_index], pr[:committed_offset]}
+  end
+
   @spec get_member_info(
           list(%{client_host: binary, member_assignment: binary, member_id: binary})
         ) ::
@@ -109,8 +103,8 @@ defmodule KafkaexLagExporter.KafkaUtils do
                            member_assignment: member_assignment,
                            member_id: consumer_id
                          } ->
-      topic_names = KafkaexLagExporter.TopicNameParser.parse_topic_names(member_assignment)
-      {topic_names, consumer_id, member_host}
+      topics = KafkaexLagExporter.TopicNameParser.parse_topic_names(member_assignment)
+      {topics, consumer_id, member_host}
     end)
   end
 end