fix mock py2/3 compatible

Yelp · Aug 30, 2016 · dc32d5c · dc32d5c
1 parent ce3e38b
commit dc32d5c
Showing 1 changed file with 9 additions and 1 deletion.
diff --git a/tests/pattern/util_test.py b/tests/pattern/util_test.py
@@ -3,6 +3,8 @@
 from __future__ import division
 from __future__ import print_function
 
+import sys
+
 import mock
 from pyparsing import Keyword
 
@@ -14,7 +16,13 @@
 from undebt.pattern.util import trailing_whitespace
 
 
-@mock.patch('__builtin__.print')
+if sys.version_info.major == 3:
+    builtin_module_name = 'builtins'
+else:
+    builtin_module_name = '__builtin__'
+
+
+@mock.patch('{}.print'.format(builtin_module_name))
 def test_debug(mock_print):
     assert_parse(
         grammar=debug(Keyword('something')),