From 08b5923d523c99d518cc67884389211e21f6fab1 Mon Sep 17 00:00:00 2001
From: Pascal Schmid <psschmid@gmx.de>
Date: Thu, 28 Mar 2024 21:23:51 +0100
Subject: [PATCH] Simplify code

---
 .../consumer_offset_fetcher.ex                | 25 +++++++------------
 test/consumer_offset_fetcher_test.exs         | 25 +++++++++++++------
 2 files changed, 26 insertions(+), 24 deletions(-)

diff --git a/lib/kafkaex_lag_exporter/consumer_offset_fetcher.ex b/lib/kafkaex_lag_exporter/consumer_offset_fetcher.ex
index 88843f6..c44c4e9 100644
--- a/lib/kafkaex_lag_exporter/consumer_offset_fetcher.ex
+++ b/lib/kafkaex_lag_exporter/consumer_offset_fetcher.ex
@@ -1,8 +1,6 @@
 defmodule KafkaexLagExporter.ConsumerOffsetFetcher do
   @moduledoc "Calculate summarized lag for each consumer group"
 
-  require Logger
-
   alias KafkaexLagExporter.KafkaUtils
 
   # TODO fix type
@@ -14,30 +12,25 @@ defmodule KafkaexLagExporter.ConsumerOffsetFetcher do
     consumer_group_names = KafkaUtils.get_consumer_group_names(endpoint)
 
     consumer_lags =
-      KafkaUtils.topic_names_for_consumer_groups(
-        endpoint,
-        [],
-        consumer_group_names
-      )
-      |> Enum.map(fn {consumer_group, topics} ->
-        {consumer_group, get_lag_for_consumer(consumer_group, topics)}
-      end)
+      KafkaUtils.topic_names_for_consumer_groups(endpoint, [], consumer_group_names)
+      |> Enum.flat_map(&get_lag_per_topic(&1))
 
     consumer_lag_sum = get_lag_for_consumer_sum(consumer_lags)
 
     %{lags: consumer_lags, sum: consumer_lag_sum}
   end
 
-  defp get_lag_for_consumer(consumer_group, topics) do
-    topics
-    |> Enum.flat_map(fn topic ->
-      KafkaUtils.lag(topic, consumer_group, :client1)
+  defp get_lag_per_topic({consumer_group, topics}) do
+    Enum.map(topics, fn topic ->
+      lag = KafkaUtils.lag(topic, consumer_group, :client1)
+      {consumer_group, topic, lag}
     end)
   end
 
   defp get_lag_for_consumer_sum(lags_per_consumer_group) do
-    lags_per_consumer_group
-    |> Enum.map(fn {topic, lag_per_partition} -> {topic, sum_topic_lag(lag_per_partition, 0)} end)
+    Enum.map(lags_per_consumer_group, fn {consumer_group, topic, lag_per_partition} ->
+      {consumer_group, topic, sum_topic_lag(lag_per_partition, 0)}
+    end)
   end
 
   defp sum_topic_lag([], acc), do: acc
diff --git a/test/consumer_offset_fetcher_test.exs b/test/consumer_offset_fetcher_test.exs
index 9ed9c51..e16838e 100644
--- a/test/consumer_offset_fetcher_test.exs
+++ b/test/consumer_offset_fetcher_test.exs
@@ -4,7 +4,11 @@ defmodule KafkaexLagExporter.ConsumerOffsetFetcher.Test do
 
   @test_consumer_group_name1 "test_consumer_1"
   @test_consumer_group_name2 "test_consumer_2"
-  @test_lags [{0, 23}, {1, 42}, {2, 666}]
+  @test_lags1 [{0, 23}, {1, 42}, {2, 666}]
+  @test_lags2 [{0, 1}, {1, 2}, {2, 3}]
+  @test_topic1 "test_topic_1"
+  @test_topic2 "test_topic_2"
+  @test_topic3 "test_topic_3"
 
   setup do
     patch(
@@ -13,15 +17,15 @@ defmodule KafkaexLagExporter.ConsumerOffsetFetcher.Test do
       fn _ -> [@test_consumer_group_name1, @test_consumer_group_name2] end
     )
 
-    patch(KafkaexLagExporter.KafkaUtils, :lag, fn _, _, _ -> @test_lags end)
+    patch(KafkaexLagExporter.KafkaUtils, :lag, &lag(&1, &2, &3))
 
     patch(
       KafkaexLagExporter.KafkaUtils,
       :topic_names_for_consumer_groups,
       fn _, _, _ ->
         [
-          {@test_consumer_group_name1, ["test_topic_1", "test_topic_2"]},
-          {@test_consumer_group_name2, ["test_topic_3"]}
+          {@test_consumer_group_name1, [@test_topic1, @test_topic2]},
+          {@test_consumer_group_name2, [@test_topic3]}
         ]
       end
     )
@@ -35,13 +39,18 @@ defmodule KafkaexLagExporter.ConsumerOffsetFetcher.Test do
     %{sum: sum, lags: lags} = KafkaexLagExporter.ConsumerOffsetFetcher.get(test_endpoint)
 
     assert sum == [
-             {@test_consumer_group_name1, 1462},
-             {@test_consumer_group_name2, 731}
+             {@test_consumer_group_name1, @test_topic1, 731},
+             {@test_consumer_group_name1, @test_topic2, 6},
+             {@test_consumer_group_name2, @test_topic3, 6}
            ]
 
     assert lags == [
-             {@test_consumer_group_name1, @test_lags ++ @test_lags},
-             {@test_consumer_group_name2, @test_lags}
+             {@test_consumer_group_name1, @test_topic1, @test_lags1},
+             {@test_consumer_group_name1, @test_topic2, @test_lags2},
+             {@test_consumer_group_name2, @test_topic3, @test_lags2}
            ]
   end
+
+  defp lag(@test_topic1, _, _), do: @test_lags1
+  defp lag(_, _, _), do: @test_lags2
 end