Add a test for conditional print + print collector.

zopefoundation · Nov 24, 2016 · aa21c35 · aa21c35
1 parent e54d9ef
commit aa21c35
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 0 deletions.
diff --git a/tests/test_print_function.py b/tests/test_print_function.py
@@ -340,3 +340,21 @@ def test_print_function_pass_print_function():
 
     ret = glb['main']()
     assert ret == '1\n2\n'
+
+
+CONDITIONAL_PRINT = """
+from __future__ import print_function
+def func(cond):
+    if cond:
+        print(1)
+    return printed
+"""
+
+
+def test_print_function_conditional_print():
+    code, errors = compiler(CONDITIONAL_PRINT)[:2]
+    glb = {'_print_': PrintCollector, '_getattr_': None}
+    six.exec_(code, glb)
+
+    assert glb['func'](True) == '1\n'
+    assert glb['func'](False) == ''
diff --git a/tests/test_print_stmt.py b/tests/test_print_stmt.py
@@ -260,3 +260,21 @@ def test_print_stmt_no_new_scope(compiler):
 
     ret = glb['class_scope']()
     assert ret == 'a\n'
+
+
+CONDITIONAL_PRINT = """
+def func(cond):
+    if cond:
+        print 1
+    return printed
+"""
+
+
+@pytest.mark.parametrize(*compilers)
+def test_print_stmt_conditional_print(compiler):
+    code, errors = compiler(CONDITIONAL_PRINT)[:2]
+    glb = {'_print_': PrintCollector, '_getattr_': None}
+    six.exec_(code, glb)
+
+    assert glb['func'](True) == '1\n'
+    assert glb['func'](False) == ''