Repository: jamescooke/flake8-aaa

Issue:
Act blocks within context managers are not found, even when marked with #act
```
    with mock.patch.object(sync_metadata, 'delay') as m_delay:
        result = url(project_biz)  # act
```

Discussion:

Patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index e45b27d..78198a4 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -14,6 +14,12 @@ Unreleased_
 See also `latest documentation
 <https://flake8-aaa.readthedocs.io/en/latest/>`_.
 
+Changed
+-------
+
+* Improved loading of Act blocks so that they can be found within context
+  managers.
+
 0.3.0_ - 2018/06/28
 -------------------
 
diff --git a/flake8_aaa/act_block.py b/flake8_aaa/act_block.py
index 07f1290..04558ed 100644
--- a/flake8_aaa/act_block.py
+++ b/flake8_aaa/act_block.py
@@ -1,4 +1,5 @@
-from .exceptions import NotActionBlock
+import ast
+
 from .helpers import node_is_pytest_raises, node_is_result_assignment
 from .types import ActBlockType
 
@@ -19,6 +20,20 @@ class ActBlock(object):
         self.node = node
         self.block_type = block_type
 
+    @classmethod
+    def build_body(cls, body):
+        """
+        Args:
+            body (list (ast.node)): List of nodes from a block.
+
+        Returns:
+            list (ActBlock)
+        """
+        act_blocks = []
+        for child_node in body:
+            act_blocks += ActBlock.build(child_node)
+        return act_blocks
+
     @classmethod
     def build(cls, node):
         """
@@ -26,18 +41,19 @@ class ActBlock(object):
             node (ast.node): A node, decorated with ``ASTTokens``.
 
         Returns:
-            ActBlock
-
-        Raises:
-            NotActionBlock: When ``node`` does not look like an Act block.
+            list(ActBlock)
         """
         if node_is_result_assignment(node):
-            return cls(node, ActBlockType.result_assignment)
+            return [cls(node, ActBlockType.result_assignment)]
         elif node_is_pytest_raises(node):
-            return cls(node, ActBlockType.pytest_raises)
+            return [cls(node, ActBlockType.pytest_raises)]
 
         # Check if line marked with '# act'
         if node.first_token.line.strip().endswith('# act'):
-            return cls(node, ActBlockType.marked_act)
+            return [cls(node, ActBlockType.marked_act)]
+
+        # Recurse if it's a context manager
+        if isinstance(node, ast.With):
+            return cls.build_body(node.body)
 
-        raise NotActionBlock()
+        return []
diff --git a/flake8_aaa/exceptions.py b/flake8_aaa/exceptions.py
index 0484c00..b0fcc95 100644
--- a/flake8_aaa/exceptions.py
+++ b/flake8_aaa/exceptions.py
@@ -2,13 +2,6 @@ class Flake8AAAException(Exception):
     pass
 
 
-class NotActionBlock(Flake8AAAException):
-    """
-    Used when parsing if lines of a function should be considered Action
-    blocks.
-    """
-
-
 class ValidationError(Flake8AAAException):
     """
     Attributes:
diff --git a/flake8_aaa/function.py b/flake8_aaa/function.py
index 6f9d757..184b5f0 100644
--- a/flake8_aaa/function.py
+++ b/flake8_aaa/function.py
@@ -1,7 +1,7 @@
 from .act_block import ActBlock
 from .arrange_block import ArrangeBlock
 from .assert_block import AssertBlock
-from .exceptions import NotActionBlock, ValidationError
+from .exceptions import ValidationError
 from .helpers import function_is_noop
 from .types import ActBlockType
 
@@ -50,12 +50,7 @@ class Function(object):
         Raises:
             ValidationError
         """
-        act_blocks = []
-        for child_node in self.node.body:
-            try:
-                act_blocks.append(ActBlock.build(child_node))
-            except NotActionBlock:
-                continue
+        act_blocks = ActBlock.build_body(self.node.body)
 
         if len(act_blocks) < 1:
             raise ValidationError(self.node.lineno, self.node.col_offset, 'AAA01 no Act block found in test')


Test patch:
diff --git a/tests/act_block/test_build.py b/tests/act_block/test_build.py
index 500980e..a704be6 100644
--- a/tests/act_block/test_build.py
+++ b/tests/act_block/test_build.py
@@ -1,7 +1,6 @@
 import pytest
 
 from flake8_aaa.act_block import ActBlock
-from flake8_aaa.exceptions import NotActionBlock
 from flake8_aaa.types import ActBlockType
 
 # TODO act blocks need testing with 'result =' indented
@@ -20,9 +19,11 @@ def test_not_actions(first_node_with_tokens):
 def test_raises_block(first_node_with_tokens):
     result = ActBlock.build(first_node_with_tokens.body[0])
 
-    assert isinstance(result, ActBlock)
-    assert result.node == first_node_with_tokens.body[0]
-    assert result.block_type == ActBlockType.pytest_raises
+    assert isinstance(result, list)
+    assert len(result) == 1
+    assert isinstance(result[0], ActBlock)
+    assert result[0].node == first_node_with_tokens.body[0]
+    assert result[0].block_type == ActBlockType.pytest_raises
 
 
 @pytest.mark.parametrize(
@@ -35,9 +36,26 @@ def test_raises_block(first_node_with_tokens):
 def test(expected_type, first_node_with_tokens):
     result = ActBlock.build(first_node_with_tokens)
 
-    assert isinstance(result, ActBlock)
-    assert result.node == first_node_with_tokens
-    assert result.block_type == expected_type
+    assert isinstance(result, list)
+    assert len(result) == 1
+    assert isinstance(result[0], ActBlock)
+    assert result[0].node == first_node_with_tokens
+    assert result[0].block_type == expected_type
+
+
+@pytest.mark.parametrize(
+    'code_str', [
+        "with mock.patch('stats.deletion_manager.deleted'):\n    result = existing_user.delete()",
+    ]
+)
+def test_nested(first_node_with_tokens):
+    result = ActBlock.build(first_node_with_tokens)
+
+    assert isinstance(result, list)
+    assert len(result) == 1
+    assert isinstance(result[0], ActBlock)
+    assert result[0].block_type == ActBlockType.result_assignment
+    assert result[0].node == first_node_with_tokens.body[0]
 
 
 @pytest.mark.parametrize(
@@ -51,5 +69,6 @@ def test(expected_type, first_node_with_tokens):
     ]
 )
 def test_not_actions(first_node_with_tokens):
-    with pytest.raises(NotActionBlock):
-        ActBlock.build(first_node_with_tokens)
+    result = ActBlock.build(first_node_with_tokens)
+
+    assert result == []
diff --git a/tests/function/test_load_act_block.py b/tests/function/test_load_act_block.py
index 38c2453..9e0d91f 100644
--- a/tests/function/test_load_act_block.py
+++ b/tests/function/test_load_act_block.py
@@ -54,6 +54,69 @@ def test_raises_in_assert(function):
     assert result.node.first_token.line == '    result = existing_user.delete()\n'
 
 
+@pytest.mark.parametrize(
+    'code_str',
+    [
+        '''
+def test(existing_user):
+    with mock.patch('stats.deletion_manager.deleted'):
+        result = existing_user.delete()
+
+    assert result is True
+    assert result.retrieved is False
+'''
+    ],
+    ids=['act in context manager'],
+)
+def test_in_cm(function):
+    result = function.load_act_block()
+
+    assert isinstance(result, ActBlock)
+    assert result.block_type == ActBlockType.result_assignment
+    assert result.node.first_token.line == '        result = existing_user.delete()\n'
+
+
+@pytest.mark.parametrize(
+    'code_str',
+    [
+        '''
+def test_no_recreate(existing_user):
+    with mock.patch('stats.creation_manager.created'):
+        with pytest.raises(ValidationError):
+            existing_user.create()
+'''
+    ],
+    ids=['pytest raises in context manager'],
+)
+def test_raises_in_cm(function):
+    result = function.load_act_block()
+
+    assert isinstance(result, ActBlock)
+    assert result.block_type == ActBlockType.pytest_raises
+    assert result.node.first_token.line == '        with pytest.raises(ValidationError):\n'
+
+
+@pytest.mark.parametrize(
+    'code_str',
+    [
+        '''
+def test_creation(stub_user):
+    with mock.patch('stats.creation_manager.created'):
+        stub_user.create()  # act
+
+    assert stub_user.exists()
+'''
+    ],
+    ids=['marked act block in context manager'],
+)
+def test_marked_in_cm(function):
+    result = function.load_act_block()
+
+    assert isinstance(result, ActBlock)
+    assert result.block_type == ActBlockType.marked_act
+    assert result.node.first_token.line == '        stub_user.create()  # act\n'
+
+
 # --- FAILURES ---
 
 
@@ -86,6 +149,17 @@ def test():
     eggs = 1  # act
 
     assert chickens + eggs == 2
+    ''',
+        '''
+def test_read(self):
+    with open('data') as data_file:
+        result = data_file.read()
+
+        assert result == ''
+
+        result = data_file.read()
+
+        assert result == ''
     ''',
     ]
 )