From a0dbf2ab995780140ce335fb96e0f797c661764b Mon Sep 17 00:00:00 2001
From: danielx123 <danielx@umich.edu>
Date: Mon, 22 Apr 2019 17:58:07 -0400
Subject: [PATCH] Adding test cases for unrolling an iterable #5062

---
 testing/test_assertrewrite.py | 47 +++++++++++++++++++++++++++++++++++
 1 file changed, 47 insertions(+)

diff --git a/testing/test_assertrewrite.py b/testing/test_assertrewrite.py
index c22ae7c6e..8add65e0a 100644
--- a/testing/test_assertrewrite.py
+++ b/testing/test_assertrewrite.py
@@ -677,6 +677,53 @@ class TestAssertionRewrite(object):
         assert "UnicodeDecodeError" not in msg
         assert "UnicodeEncodeError" not in msg
 
+    def test_generator(self, testdir):
+        testdir.makepyfile(
+            """
+            def check_even(num):
+                if num % 2 == 0:
+                    return True
+                return False
+
+            def test_generator():
+                odd_list = list(range(1,9,2))
+                assert all(check_even(num) for num in odd_list)"""
+        )
+        result = testdir.runpytest()
+        result.stdout.fnmatch_lines(["*assert False*", "*where False = check_even(1)*"])
+
+    def test_list_comprehension(self, testdir):
+        testdir.makepyfile(
+            """
+            def check_even(num):
+                if num % 2 == 0:
+                    return True
+                return False
+
+            def test_list_comprehension():
+                odd_list = list(range(1,9,2))
+                assert all([check_even(num) for num in odd_list])"""
+        )
+        result = testdir.runpytest()
+        result.stdout.fnmatch_lines(["*assert False*", "*where False = check_even(1)*"])
+
+    def test_for_loop(self, testdir):
+        testdir.makepyfile(
+            """
+            def check_even(num):
+                if num % 2 == 0:
+                    return True
+                return False
+
+            def test_for_loop():
+                odd_list = list(range(1,9,2))
+                for num in odd_list:
+                    assert check_even(num)
+        """
+        )
+        result = testdir.runpytest()
+        result.stdout.fnmatch_lines(["*assert False*", "*where False = check_even(1)*"])
+
 
 class TestRewriteOnImport(object):
     def test_pycache_is_a_file(self, testdir):