Patch print function instead of sys.stdout

1 files changed, 12 insertions, 9 deletions

diff --git a/tests/unit/test_base_formatter.py b/tests/unit/test_base_formatter.py
index 9da8c65..1cf7562 100644
--- a/tests/unit/test_base_formatter.py
+++ b/tests/unit/test_base_formatter.py
@@ -1,5 +1,4 @@
 """Tests for the BaseFormatter object."""
-import io
 import optparse
 
 import mock
@@ -81,19 +80,23 @@ def test_show_source_updates_physical_line_appropriately(line, column):
 @pytest.mark.parametrize('tee', [False, True])
 def test_write_uses_an_output_file(tee):
     """Verify that we use the output file when it's present."""
-    line = u'Something to write'
-    source = u'source'
+    line = 'Something to write'
+    source = 'source'
     filemock = mock.Mock()
 
-    with mock.patch('sys.stdout', new_callable=io.StringIO) as stdout:
-        formatter = base.BaseFormatter(options(tee=tee))
-        formatter.output_fd = filemock
+    formatter = base.BaseFormatter(options(tee=tee))
+    formatter.output_fd = filemock
+
+    with mock.patch('flake8.formatting.base.print') as print_func:
         formatter.write(line, source)
         if tee:
-            output = line + formatter.newline + source + formatter.newline
+            assert print_func.called
+            assert print_func.mock_calls == [
+                mock.call(line),
+                mock.call(source),
+            ]
         else:
-            output = ''
-        assert stdout.getvalue() == output
+            assert not print_func.called
 
     assert filemock.write.called is True
     assert filemock.write.call_count == 2