More test repairs.

1 files changed, 9 insertions, 3 deletions

diff --git a/tests/unittests/test_datasource/test_smartos.py b/tests/unittests/test_datasource/test_smartos.py
index 01b9b73e..2fb9e1b6 100644
--- a/tests/unittests/test_datasource/test_smartos.py
+++ b/tests/unittests/test_datasource/test_smartos.py
@@ -36,6 +36,12 @@ import tempfile
 import stat
 import uuid
 
+def b64(source):
+    # In Python 3, b64encode only accepts bytes and returns bytes.
+    if not isinstance(source, bytes):
+        source = source.encode('utf-8')
+    return base64.b64encode(source).decode('us-ascii')
+
 
 MOCK_RETURNS = {
     'hostname': 'test-host',
@@ -233,7 +239,7 @@ class TestSmartOSDataSource(helpers.FilesystemMockingTestCase):
         my_returns = MOCK_RETURNS.copy()
         my_returns['base64_all'] = "true"
         for k in ('hostname', 'cloud-init:user-data'):
-            my_returns[k] = base64.b64encode(my_returns[k])
+            my_returns[k] = b64(my_returns[k])
 
         dsrc = self._get_ds(mockdata=my_returns)
         ret = dsrc.get_data()
@@ -254,7 +260,7 @@ class TestSmartOSDataSource(helpers.FilesystemMockingTestCase):
         my_returns['b64-cloud-init:user-data'] = "true"
         my_returns['b64-hostname'] = "true"
         for k in ('hostname', 'cloud-init:user-data'):
-            my_returns[k] = base64.b64encode(my_returns[k])
+            my_returns[k] = b64(my_returns[k])
 
         dsrc = self._get_ds(mockdata=my_returns)
         ret = dsrc.get_data()
@@ -270,7 +276,7 @@ class TestSmartOSDataSource(helpers.FilesystemMockingTestCase):
         my_returns = MOCK_RETURNS.copy()
         my_returns['base64_keys'] = 'hostname,ignored'
         for k in ('hostname',):
-            my_returns[k] = base64.b64encode(my_returns[k])
+            my_returns[k] = b64(my_returns[k])
 
         dsrc = self._get_ds(mockdata=my_returns)
         ret = dsrc.get_data()