bpo-32933: Implement __iter__ method on mock_open() (GH-5974)

diff --git a/Doc/library/unittest.mock.rst b/Doc/library/unittest.mock.rst
index fd4e067546e5ea6c53d5dcf834964c8c67b9d51e..d1b18d08f797b4bb28bb8d2ca73907653fed0b24 100644 (file)
--- a/Doc/library/unittest.mock.rst
+++ b/Doc/library/unittest.mock.rst
@@ -2095,6 +2095,10 @@ mock_open
    .. versionchanged:: 3.5
       *read_data* is now reset on each call to the *mock*.
 
+   .. versionchanged:: 3.8
+      Added :meth:`__iter__` to implementation so that iteration (such as in for
+      loops) correctly consumes *read_data*.
+
 Using :func:`open` as a context manager is a great way to ensure your file handles
 are closed properly and is becoming common::
 

diff --git a/Lib/unittest/mock.py b/Lib/unittest/mock.py
index db1e642c00b7f9c8b86eb3c81b880b35b9851b33..83026e6f3bd398fe082237340b4355d37f212ea0 100644 (file)
--- a/Lib/unittest/mock.py
+++ b/Lib/unittest/mock.py
@@ -2358,14 +2358,16 @@ def mock_open(mock=None, read_data=''):
         return type(read_data)().join(_state[0])
 
     def _readline_side_effect():
+        yield from _iter_side_effect()
+        while True:
+            yield type(read_data)()
+
+    def _iter_side_effect():
         if handle.readline.return_value is not None:
             while True:
                 yield handle.readline.return_value
         for line in _state[0]:
             yield line
-        while True:
-            yield type(read_data)()
-
 
     global file_spec
     if file_spec is None:
@@ -2389,6 +2391,7 @@ def mock_open(mock=None, read_data=''):
     _state[1] = _readline_side_effect()
     handle.readline.side_effect = _state[1]
     handle.readlines.side_effect = _readlines_side_effect
+    handle.__iter__.side_effect = _iter_side_effect
 
     def reset_data(*args, **kwargs):
         _state[0] = _iterate_read_data(read_data)

diff --git a/Lib/unittest/test/testmock/testmock.py b/Lib/unittest/test/testmock/testmock.py
index b64c8663d212266e7d9e40697b50ea52d8bcd845..c7bfa277b511ca96aae59072f146b0bb79b528ef 100644 (file)
--- a/Lib/unittest/test/testmock/testmock.py
+++ b/Lib/unittest/test/testmock/testmock.py
@@ -1450,6 +1450,16 @@ class MockTest(unittest.TestCase):
         f2_data = f2.read()
         self.assertEqual(f1_data, f2_data)
 
+    def test_mock_open_dunder_iter_issue(self):
+        # Test dunder_iter method generates the expected result and
+        # consumes the iterator.
+        mocked_open = mock.mock_open(read_data='Remarkable\nNorwegian Blue')
+        f1 = mocked_open('a-name')
+        lines = [line for line in f1]
+        self.assertEqual(lines[0], 'Remarkable\n')
+        self.assertEqual(lines[1], 'Norwegian Blue')
+        self.assertEqual(list(f1), [])
+
     def test_mock_open_write(self):
         # Test exception in file writing write()
         mock_namedtemp = mock.mock_open(mock.MagicMock(name='JLV'))

diff --git a/Lib/unittest/test/testmock/testwith.py b/Lib/unittest/test/testmock/testwith.py
index a7bee730030190e4ed456f817369aeea0e409296..43b36a1199526122d2e252385fb35bac827f23a7 100644 (file)
--- a/Lib/unittest/test/testmock/testwith.py
+++ b/Lib/unittest/test/testmock/testwith.py
@@ -188,6 +188,7 @@ class TestMockOpen(unittest.TestCase):
 
     def test_readline_data(self):
         # Check that readline will return all the lines from the fake file
+        # And that once fully consumed, readline will return an empty string.
         mock = mock_open(read_data='foo\nbar\nbaz\n')
         with patch('%s.open' % __name__, mock, create=True):
             h = open('bar')
@@ -197,6 +198,7 @@ class TestMockOpen(unittest.TestCase):
         self.assertEqual(line1, 'foo\n')
         self.assertEqual(line2, 'bar\n')
         self.assertEqual(line3, 'baz\n')
+        self.assertEqual(h.readline(), '')
 
         # Check that we properly emulate a file that doesn't end in a newline
         mock = mock_open(read_data='foo')
@@ -204,6 +206,19 @@ class TestMockOpen(unittest.TestCase):
             h = open('bar')
             result = h.readline()
         self.assertEqual(result, 'foo')
+        self.assertEqual(h.readline(), '')
+
+
+    def test_dunder_iter_data(self):
+        # Check that dunder_iter will return all the lines from the fake file.
+        mock = mock_open(read_data='foo\nbar\nbaz\n')
+        with patch('%s.open' % __name__, mock, create=True):
+            h = open('bar')
+            lines = [l for l in h]
+        self.assertEqual(lines[0], 'foo\n')
+        self.assertEqual(lines[1], 'bar\n')
+        self.assertEqual(lines[2], 'baz\n')
+        self.assertEqual(h.readline(), '')
 
 
     def test_readlines_data(self):

diff --git a/Misc/NEWS.d/next/Library/2018-04-30-22-43-31.bpo-32933.M3iI_y.rst b/Misc/NEWS.d/next/Library/2018-04-30-22-43-31.bpo-32933.M3iI_y.rst
new file mode 100644 (file)
index 0000000..4de7a8f
--- /dev/null
+++ b/Misc/NEWS.d/next/Library/2018-04-30-22-43-31.bpo-32933.M3iI_y.rst
@@ -0,0 +1,2 @@
+:func:`unittest.mock.mock_open` now supports iteration over the file
+contents. Patch by Tony Flury.