add ensure_binary/str/text helper functions (#204)

1 files changed, 58 insertions, 0 deletions

diff --git a/test_six.py b/test_six.py
index 872b087..980cdf3 100644
--- a/test_six.py
+++ b/test_six.py
@@ -932,3 +932,61 @@ def test_python_2_unicode_compatible():
         assert str(my_test) == six.u("hello")
 
     assert getattr(six.moves.builtins, 'bytes', str)(my_test) == six.b("hello")
+
+
+class EnsureTests:
+
+    # grinning face emoji
+    UNICODE_EMOJI = six.u("\U0001F600")
+    BINARY_EMOJI = b"\xf0\x9f\x98\x80"
+
+    def test_ensure_binary_raise_type_error(self):
+        with py.test.raises(TypeError):
+            six.ensure_str(8)
+
+    def test_errors_and_encoding(self):
+        six.ensure_binary(self.UNICODE_EMOJI, encoding='latin-1', errors='ignore')
+        with py.test.raises(UnicodeEncodeError):
+            six.ensure_binary(self.UNICODE_EMOJI, encoding='latin-1', errors='strict')
+
+    def test_ensure_binary_raise(self):
+        converted_unicode = six.ensure_binary(self.UNICODE_EMOJI, encoding='utf-8', errors='strict')
+        converted_binary = six.ensure_binary(self.BINARY_EMOJI, encoding="utf-8", errors='strict')
+        if six.PY2:
+            # PY2: unicode -> str
+            assert converted_unicode == self.BINARY_EMOJI and isinstance(converted_unicode, str)
+            # PY2: str -> str
+            assert converted_binary == self.BINARY_EMOJI and isinstance(converted_binary, str)
+        else:
+            # PY3: str -> bytes
+            assert converted_unicode == self.BINARY_EMOJI and isinstance(converted_unicode, bytes)
+            # PY3: bytes -> bytes
+            assert converted_binary == self.BINARY_EMOJI and isinstance(converted_binary, bytes)
+
+    def test_ensure_str(self):
+        converted_unicode = six.ensure_str(self.UNICODE_EMOJI, encoding='utf-8', errors='strict')
+        converted_binary = six.ensure_str(self.BINARY_EMOJI, encoding="utf-8", errors='strict')
+        if six.PY2:
+            # PY2: unicode -> str
+            assert converted_unicode == self.BINARY_EMOJI and isinstance(converted_unicode, str)
+            # PY2: str -> str
+            assert converted_binary == self.BINARY_EMOJI and isinstance(converted_binary, str)
+        else:
+            # PY3: str -> str
+            assert converted_unicode == self.UNICODE_EMOJI and isinstance(converted_unicode, str)
+            # PY3: bytes -> str
+            assert converted_binary == self.UNICODE_EMOJI and isinstance(converted_unicode, str)
+
+    def test_ensure_text(self):
+        converted_unicode = six.ensure_text(self.UNICODE_EMOJI, encoding='utf-8', errors='strict')
+        converted_binary = six.ensure_text(self.BINARY_EMOJI, encoding="utf-8", errors='strict')
+        if six.PY2:
+            # PY2: unicode -> unicode
+            assert converted_unicode == self.UNICODE_EMOJI and isinstance(converted_unicode, unicode)
+            # PY2: str -> unicode
+            assert converted_binary == self.UNICODE_EMOJI and isinstance(converted_unicode, unicode)
+        else:
+            # PY3: str -> str
+            assert converted_unicode == self.UNICODE_EMOJI and isinstance(converted_unicode, str)
+            # PY3: bytes -> str
+            assert converted_binary == self.UNICODE_EMOJI and isinstance(converted_unicode, str)