add tests for python_2_unicode_comapatible

2 files changed, 22 insertions, 1 deletions

diff --git a/six.py b/six.py
index 35530a3..e7667af 100644
--- a/six.py
+++ b/six.py
@@ -746,7 +746,7 @@ def python_2_unicode_compatible(klass):
     To support Python 2 and 3 with a single code base, define a __str__ method
     returning text and apply this decorator to the class.
     """
-    if six.PY2:
+    if PY2:
         if '__str__' not in klass.__dict__:
             raise ValueError("@python_2_unicode_compatible cannot be applied "
                              "to %s because it doesn't define __str__()." %
diff --git a/test_six.py b/test_six.py
index 0125d6b..c744b29 100644
--- a/test_six.py
+++ b/test_six.py
@@ -734,3 +734,24 @@ def test_add_metaclass():
         __slots__ = "__weakref__",
     MySlotsWeakref = six.add_metaclass(Meta)(MySlotsWeakref)
     assert type(MySlotsWeakref) is Meta
+
+
+def test_python_2_unicode_compatible():
+    @six.python_2_unicode_compatible
+    class MyTest(object):
+        def __str__(self):
+            return six.u('hello')
+
+        def __bytes__(self):
+            return six.b('hello')
+
+    my_test = MyTest()
+
+    if six.PY2:
+        assert str(my_test) == six.b("hello")
+        assert unicode(my_test) == six.u("hello")
+    elif six.PY3:
+        assert bytes(my_test) == six.b("hello")
+        assert str(my_test) == six.u("hello")
+
+    assert getattr(six.moves.builtins, 'bytes', str)(my_test) == six.b("hello")