From 2ca6b859c86061b0759fae311a871009961fb494 Mon Sep 17 00:00:00 2001
From: JessicaBell00 <110268278+JessicaBell00@users.noreply.github.com>
Date: Mon, 2 Sep 2024 13:45:16 +1200
Subject: [PATCH] Added checking for raises in 'for' branches

Also added relevant tests
---
 .../pylint_guidelines_checker.py              | 16 +++---
 .../tests/test_pylint_custom_plugins.py       | 52 +++++++++++++++++++
 2 files changed, 60 insertions(+), 8 deletions(-)

diff --git a/tools/pylint-extensions/azure-pylint-guidelines-checker/pylint_guidelines_checker.py b/tools/pylint-extensions/azure-pylint-guidelines-checker/pylint_guidelines_checker.py
index 44cd8660437..35870d69bdb 100644
--- a/tools/pylint-extensions/azure-pylint-guidelines-checker/pylint_guidelines_checker.py
+++ b/tools/pylint-extensions/azure-pylint-guidelines-checker/pylint_guidelines_checker.py
@@ -2741,19 +2741,19 @@ def check_for_raise(self, node):
             if isinstance(i, astroid.Raise):
                 self.check_for_logging(node)
             # Check for any nested 'If' branches
-            if isinstance(i, astroid.If):
+            if isinstance(i, astroid.If) or isinstance(i, astroid.For):
                 self.check_for_raise(i.body)
-                # Check any 'elif' and 'else' branches
-                elif_statements = i.orelse
-                while len(elif_statements) == 1:
-                    if isinstance(elif_statements[0], astroid.If):
-                        for x in elif_statements:
+                # Check any 'elif', 'else' or 'for' branches
+                orelse_statements = i.orelse
+                while len(orelse_statements) == 1:
+                    if isinstance(orelse_statements[0], astroid.If):
+                        for x in orelse_statements:
                             self.check_for_raise(x.body)
-                            elif_statements = x.orelse
+                            orelse_statements = x.orelse
                     else:
                         break
                 # Check 'else' body for raise
-                self.check_for_raise(elif_statements)
+                self.check_for_raise(orelse_statements)
 
     def check_for_logging(self, node):
         """ Helper function - checks 'Expr' nodes to see if logging has occurred at 'warning' or 'error'
diff --git a/tools/pylint-extensions/azure-pylint-guidelines-checker/tests/test_pylint_custom_plugins.py b/tools/pylint-extensions/azure-pylint-guidelines-checker/tests/test_pylint_custom_plugins.py
index 8281b25fbbb..54922591d74 100644
--- a/tools/pylint-extensions/azure-pylint-guidelines-checker/tests/test_pylint_custom_plugins.py
+++ b/tools/pylint-extensions/azure-pylint-guidelines-checker/tests/test_pylint_custom_plugins.py
@@ -5108,3 +5108,55 @@ def test_extra_nested_branches_exception_logged(self):
                 )
         ):
             self.checker.visit_try(try_node)
+
+    def test_for_and_nested_for_branches_exception_logged(self):
+        try_node, expression_node_a, expression_node_b = astroid.extract_node(
+            '''
+                try: #@
+                    add = 1 + 2
+                except Exception as e:
+                    y = 3
+                    for x in y:
+                        logging.error(f"F: {e}") #@
+                        raise SystemError("Uh oh!") from e
+                    if e.code == "Retryable":
+                        for z in y:
+                            logging.error(f"F: {e}") #@
+                            raise SystemError("Uh oh!") from e
+            ''')
+        with self.assertAddsMessages(
+                pylint.testutils.MessageTest(
+                    msg_id="do-not-log-raised-errors",
+                    line=7,
+                    node=expression_node_a.parent,
+                    col_offset=8,
+                    end_line=7,
+                    end_col_offset=32,
+                ),
+                pylint.testutils.MessageTest(
+                    msg_id="do-not-log-raised-errors",
+                    line=11,
+                    node=expression_node_b.parent,
+                    col_offset=12,
+                    end_line=11,
+                    end_col_offset=36,
+                )
+        ):
+            self.checker.visit_try(try_node)
+
+    def test_for_and_nested_for_branches_exception_not_logged(self):
+        try_node, expression_node_a, expression_node_b = astroid.extract_node(
+            '''
+                try: #@
+                    add = 1 + 2
+                except Exception as e:
+                    y = 3
+                    for x in y:
+                        logging.error(f"F: {e}") #@
+                    if e.code == "Retryable":
+                        for z in y:
+                            logging.error(f"F: {e}") #@
+            ''')
+        with self.assertNoMessages():
+            self.checker.visit_try(try_node)
+