[kafka_consumer] Fix context count for max_partition_contexts (DataDog#21223)

* Fix context count

# Conflicts:
#	kafka_consumer/tests/test_unit.py

* Edit changelog.

* Fix changelog PR number.

* Run the formatter.

Author: nubtron

kafka_consumer/changelog.d/21223.fixed

@@ -0,0 +1 @@
+Fix undercount of contexts for the max_partition_contexts configuration option.

kafka_consumer/datadog_checks/kafka_consumer/kafka_consumer.py

@@ -62,8 +62,9 @@ def check(self, _):
         broker_timestamps = defaultdict(dict)
         cluster_id = ""
         persistent_cache_key = "broker_timestamps_"
+        consumer_contexts_count = self.count_consumer_contexts(consumer_offsets)
         try:
-            if len(consumer_offsets) < self._context_limit:
+            if consumer_contexts_count < self._context_limit:
                 # Fetch highwater offsets
                 # Expected format: ({(topic, partition): offset}, cluster_id)
                 highwater_offsets, cluster_id = self.get_highwater_offsets(consumer_offsets)
@@ -80,7 +81,7 @@ def check(self, _):
                 self.client.close_admin_client()
             raise
 
-        total_contexts = sum(len(v) for v in consumer_offsets.values()) + len(highwater_offsets)
+        total_contexts = consumer_contexts_count + len(highwater_offsets)
         self.log.debug(
             "Total contexts: %s, Consumer offsets: %s, Highwater offsets: %s",
             total_contexts,
@@ -108,6 +109,9 @@ def check(self, _):
         if self.config._close_admin_client:
             self.client.close_admin_client()
 
+    def count_consumer_contexts(self, consumer_offsets):
+        return sum(len(offsets) for offsets in consumer_offsets.values())
+
     def get_consumer_offsets(self):
         # {(consumer_group, topic, partition): offset}
         self.log.debug('Getting consumer offsets')

kafka_consumer/tests/test_unit.py

@@ -1001,6 +1001,15 @@ def test_build_schema_none_handling():
         build_protobuf_schema(None)
 
 
+def test_count_consumer_contexts(check, kafka_instance):
+    kafka_consumer_check = check(kafka_instance)
+    consumer_offsets = {
+        'consumer_group1': {('topic1', 'partition0'): 1, ('topic1', 'partition1'): 2},  # 2 contexts
+        'consumer_group2': {('topic2', 'partition0'): 3},  # 1 context
+    }
+    assert kafka_consumer_check.count_consumer_contexts(consumer_offsets) == 3
+
+
 def test_consumer_group_state_fetched_once_per_group(check, kafka_instance, dd_run_check, aggregator):
     mock_client = seed_mock_client()
     # Set up two partitions for same topic to check multiple contexts in same consumer group