Added unittest aliases.

3 files changed, 88 insertions, 0 deletions

diff --git a/documentation/index.rst b/documentation/index.rst
index 7851421..6b3ad1e 100644
--- a/documentation/index.rst
+++ b/documentation/index.rst
@@ -746,6 +746,39 @@ Contains classes from Python 3's :mod:`py3:urllib.response` and Python 2's:
 * :class:`py2:urllib.addinfourl`
 
 
+unittest assertions
+<<<<<<<<<<<<<<<<<<<
+
+.. currentmodule:: six
+
+Contains compatibility shims for unittest assertions that have been renamed.
+The parameters are the same as their aliases, but you must pass the test method
+as the first argument. For example::
+
+    import six
+    import unittest
+
+    class TestAssertCountEqual(unittest.TestCase):
+        def test(self):
+            six.assertCountEqual(self, (1, 2), [2, 1])
+
+
+.. function:: assertCountEqual()
+
+   Alias for :meth:`~py3:unittest.TestCase.assertCountEqual` on Python 3 and
+   :meth:`~py2:unittest.TestCase.assertItemsEqual` on Python 2.
+
+.. function:: assertRaisesRegex()
+
+   Alias for :meth:`~py3:unittest.TestCase.assertRaisesRegex` on Python 3 and
+   :meth:`~py2:unittest.TestCase.assertRaisesRegexp` on Python 2.
+
+.. function:: assertRegex()
+
+   Alias for :meth:`~py3:unittest.TestCase.assertRegex` on Python 3 and
+   :meth:`~py2:unittest.TestCase.assertRegexpMatches` on Python 2.
+
+
 Advanced - Customizing renames
 <<<<<<<<<<<<<<<<<<<<<<<<<<<<<<
 
diff --git a/six.py b/six.py
index a8cdb22..303fd2d 100644
--- a/six.py
+++ b/six.py
@@ -610,6 +610,9 @@ if PY3:
     import io
     StringIO = io.StringIO
     BytesIO = io.BytesIO
+    _assertCountEqual = "assertCountEqual"
+    _assertRaisesRegex = "assertRaisesRegex"
+    _assertRegex = "assertRegex"
 else:
     def b(s):
         return s
@@ -625,10 +628,25 @@ else:
     iterbytes = functools.partial(itertools.imap, ord)
     import StringIO
     StringIO = BytesIO = StringIO.StringIO
+    _assertCountEqual = "assertItemsEqual"
+    _assertRaisesRegex = "assertRaisesRegexp"
+    _assertRegex = "assertRegexpMatches"
 _add_doc(b, """Byte literal""")
 _add_doc(u, """Text literal""")
 
 
+def assertCountEqual(self, *args, **kwargs):
+    return getattr(self, _assertCountEqual)(*args, **kwargs)
+
+
+def assertRaisesRegex(self, *args, **kwargs):
+    return getattr(self, _assertRaisesRegex)(*args, **kwargs)
+
+
+def assertRegex(self, *args, **kwargs):
+    return getattr(self, _assertRegex)(*args, **kwargs)
+
+
 if PY3:
     exec_ = getattr(moves.builtins, "exec")
 
diff --git a/test_six.py b/test_six.py
index b0ccd8d..c023862 100644
--- a/test_six.py
+++ b/test_six.py
@@ -1,6 +1,7 @@
 import operator
 import sys
 import types
+import unittest
 
 import py
 
@@ -773,3 +774,39 @@ def test_add_metaclass():
         __slots__ = "__weakref__",
     MySlotsWeakref = six.add_metaclass(Meta)(MySlotsWeakref)
     assert type(MySlotsWeakref) is Meta
+
+
+@py.test.mark.skipif("sys.version_info[:2] < (2, 7)")
+def test_assertCountEqual():
+    class TestAssertCountEqual(unittest.TestCase):
+        def test(self):
+            with self.assertRaises(AssertionError):
+                six.assertCountEqual(self, (1, 2), [3, 4, 5])
+
+            six.assertCountEqual(self, (1, 2), [2, 1])
+
+    TestAssertCountEqual('test').test()
+
+
+def test_assertRegex():
+    class TestAssertRegex(unittest.TestCase):
+        def test(self):
+            with self.assertRaises(AssertionError):
+                six.assertRegex(self, 'test', r'^a')
+
+            six.assertRegex(self, 'test', r'^t')
+
+    TestAssertRegex('test').test()
+
+
+def test_assertRaisesRegex():
+    class TestAssertRaisesRegex(unittest.TestCase):
+        def test(self):
+            with six.assertRaisesRegex(self, AssertionError, '^Foo'):
+                raise AssertionError('Foo')
+
+            with self.assertRaises(AssertionError):
+                with six.assertRaisesRegex(self, AssertionError, r'^Foo'):
+                    raise AssertionError('Bar')
+
+    TestAssertRaisesRegex('test').test()