fixup

airbnb · Oct 4, 2019 · c5b0267 · c5b0267
1 parent c567380
commit c5b0267
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 5 deletions.
diff --git a/streamalert/alert_merger/main.py b/streamalert/alert_merger/main.py
@@ -87,16 +87,15 @@ def _alert_generator(self, rule_name):
         To limit memory consumption, the generator yields a maximum number of alerts, defined
         by self._alert_generator_limit.
         """
-        alert_count = 0
-        for record in self.table.get_alert_records(rule_name, self.alert_proc_timeout):
+        generator = self.table.get_alert_records(rule_name, self.alert_proc_timeout)
+        for idx, record in enumerate(generator, start=1):
             try:
                 yield Alert.create_from_dynamo_record(record)
-                alert_count += 1
             except AlertCreationError:
                 LOGGER.exception('Invalid alert record %s', record)
                 continue
 
-            if alert_count >= self._alert_generator_limit:
+            if idx >= self._alert_generator_limit:
                 LOGGER.warning(
                     'Alert Merger reached alert limit of %d for rule "%s"',
                     self._alert_generator_limit,

diff --git a/streamalert/shared/alert_table.py b/streamalert/shared/alert_table.py
@@ -81,7 +81,7 @@ def rule_names_generator(self):
         it is possible for certain names to get skipped.
 
         Returns:
-            generator
+            Generator[str]
         """
         kwargs = {
             'ProjectionExpression': 'RuleName',