[test_unicode_literals] Fix test

2 files changed, 15 insertions, 3 deletions

diff --git a/test/helper.py b/test/helper.py
index 9a7f0746e..8a820526a 100644
--- a/test/helper.py
+++ b/test/helper.py
@@ -161,7 +161,9 @@ def assertRegexpMatches(self, text, regexp, msg=None):
     else:
         m = re.match(regexp, text)
         if not m:
-            note = 'Regexp didn\'t match: %r not found in %r' % (regexp, text)
+            note = 'Regexp didn\'t match: %r not found' % (regexp)
+            if len(text) < 1000:
+                note += ' in %r' % text
             if msg is None:
                 msg = note
             else:
diff --git a/test/test_unicode_literals.py b/test/test_unicode_literals.py
index 2cc431b0b..d3cba869b 100644
--- a/test/test_unicode_literals.py
+++ b/test/test_unicode_literals.py
@@ -1,5 +1,11 @@
 from __future__ import unicode_literals
 
+# Allow direct execution
+import os
+import sys
+import unittest
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+
 import io
 import os
 import re
@@ -14,6 +20,9 @@ IGNORED_FILES = [
 ]
 
 
+from helper import assertRegexpMatches
+
+
 class TestUnicodeLiterals(unittest.TestCase):
     def test_all_files(self):
         for dirpath, _, filenames in os.walk(rootDir):
@@ -29,9 +38,10 @@ class TestUnicodeLiterals(unittest.TestCase):
 
                 if "'" not in code and '"' not in code:
                     continue
-                self.assertRegexpMatches(
+                assertRegexpMatches(
+                    self,
                     code,
-                    r'(?:#.*\n*)?from __future__ import (?:[a-z_]+,\s*)*unicode_literals',
+                    r'(?:(?:#.*?|\s*)\n)*from __future__ import (?:[a-z_]+,\s*)*unicode_literals',
                     'unicode_literals import  missing in %s' % fn)
 
                 m = re.search(r'(?<=\s)u[\'"](?!\)|,|$)', code)