From 539aa06f6490701d1e462561b299fed8e5130f08 Mon Sep 17 00:00:00 2001
From: Mohammad Azim Khan <Azim.Khan@arm.com>
Date: Fri, 6 Jul 2018 00:29:50 +0100
Subject: [PATCH] Fix Python 2 & 3 compatibility in test_generate_test_code.py

---
 tests/scripts/test_generate_test_code.py | 73 +++++++++++++++++++-----
 1 file changed, 60 insertions(+), 13 deletions(-)

diff --git a/tests/scripts/test_generate_test_code.py b/tests/scripts/test_generate_test_code.py
index 149159c8c..2ef12e18d 100755
--- a/tests/scripts/test_generate_test_code.py
+++ b/tests/scripts/test_generate_test_code.py
@@ -23,9 +23,19 @@ Unit tests for generate_test_code.py
 """
 
 
-from StringIO import StringIO
+try:
+    # Python 2
+    from StringIO import StringIO
+except ImportError:
+    # Python 3
+    from io import StringIO
 from unittest import TestCase, main as unittest_main
-from mock import patch
+try:
+    # Python 2
+    from mock import patch
+except ImportError:
+    # Python 3
+    from unittest.mock import patch
 from generate_test_code import gen_dependencies, gen_dependencies_one_line
 from generate_test_code import gen_function_wrapper, gen_dispatch
 from generate_test_code import parse_until_pattern, GeneratorInputError
@@ -307,9 +317,16 @@ class StringIOWrapper(StringIO, object):
 
         :return: Line read from file.
         """
-        line = super(StringIOWrapper, self).next()
+        parent = super(StringIOWrapper, self)
+        if getattr(parent, 'next', None):
+            # Python 2
+            line = parent.next()
+        else:
+            # Python 3
+            line = parent.__next__()
         return line
 
+    # Python 3
     __next__ = next
 
     def readline(self, length=0):
@@ -532,6 +549,38 @@ class ParseFunctionCode(TestCase):
     Test suite for testing parse_function_code()
     """
 
+    def assert_raises_regex(self, exp, regex, func, *args):
+        """
+        Python 2 & 3 portable wrapper of assertRaisesRegex(p)? function.
+
+        :param exp: Exception type expected to be raised by cb.
+        :param regex: Expected exception message
+        :param func: callable object under test
+        :param args: variable positional arguments
+        """
+        parent = super(ParseFunctionCode, self)
+
+        # Pylint does not appreciate that the super method called
+        # conditionally can be available in other Python version
+        # then that of Pylint.
+        # Workaround is to call the method via getattr.
+        # Pylint ignores that the method got via getattr is
+        # conditionally executed. Method has to be a callable.
+        # Hence, using a dummy callable for getattr default.
+        dummy = lambda *x: None
+        # First Python 3 assertRaisesRegex is checked, since Python 2
+        # assertRaisesRegexp is also available in Python 3 but is
+        # marked deprecated.
+        for name in ('assertRaisesRegex', 'assertRaisesRegexp'):
+            method = getattr(parent, name, dummy)
+            if method is not dummy:
+                method(exp, regex, func, *args)
+                break
+        else:
+            raise AttributeError(" 'ParseFunctionCode' object has no attribute"
+                                 " 'assertRaisesRegex' or 'assertRaisesRegexp'"
+                                )
+
     def test_no_function(self):
         """
         Test no test function found.
@@ -544,8 +593,8 @@ function
 '''
         stream = StringIOWrapper('test_suite_ut.function', data)
         err_msg = 'file: test_suite_ut.function - Test functions not found!'
-        self.assertRaisesRegexp(GeneratorInputError, err_msg,
-                                parse_function_code, stream, [], [])
+        self.assert_raises_regex(GeneratorInputError, err_msg,
+                                 parse_function_code, stream, [], [])
 
     def test_no_end_case_comment(self):
         """
@@ -560,8 +609,8 @@ void test_func()
         stream = StringIOWrapper('test_suite_ut.function', data)
         err_msg = r'file: test_suite_ut.function - '\
                   'end case pattern .*? not found!'
-        self.assertRaisesRegexp(GeneratorInputError, err_msg,
-                                parse_function_code, stream, [], [])
+        self.assert_raises_regex(GeneratorInputError, err_msg,
+                                 parse_function_code, stream, [], [])
 
     @patch("generate_test_code.parse_function_arguments")
     def test_function_called(self,
@@ -678,8 +727,8 @@ exit:
         data = 'int entropy_threshold( char * a, data_t * h, int result )'
         err_msg = 'file: test_suite_ut.function - Test functions not found!'
         stream = StringIOWrapper('test_suite_ut.function', data)
-        self.assertRaisesRegexp(GeneratorInputError, err_msg,
-                                parse_function_code, stream, [], [])
+        self.assert_raises_regex(GeneratorInputError, err_msg,
+                                 parse_function_code, stream, [], [])
 
     @patch("generate_test_code.gen_dispatch")
     @patch("generate_test_code.gen_dependencies")
@@ -1155,8 +1204,7 @@ dhm_do_dhm:10:"93450983094850938450983409623":10:"9345098304850938450983409622"
             for _, _, _, _ in parse_test_data(stream):
                 pass
         except GeneratorInputError as err:
-            pass
-        self.assertEqual(type(err), GeneratorInputError)
+            self.assertEqual(type(err), GeneratorInputError)
 
     def test_incomplete_data(self):
         """
@@ -1174,8 +1222,7 @@ depends_on:YAHOO
             for _, _, _, _ in parse_test_data(stream):
                 pass
         except GeneratorInputError as err:
-            pass
-        self.assertEqual(type(err), GeneratorInputError)
+            self.assertEqual(type(err), GeneratorInputError)
 
 
 class GenDepCheck(TestCase):