Merge pull request #4 from zdc/T2117

Cloud-init: T2117: Updated to 20.1 version

1 files changed, 5 insertions, 4 deletions

diff --git a/cloudinit/cmd/tests/test_main.py b/cloudinit/cmd/tests/test_main.py
index a1e534fb..384fddc6 100644
--- a/cloudinit/cmd/tests/test_main.py
+++ b/cloudinit/cmd/tests/test_main.py
@@ -3,11 +3,12 @@
 from collections import namedtuple
 import copy
 import os
-from six import StringIO
+from io import StringIO
 
 from cloudinit.cmd import main
+from cloudinit import safeyaml
 from cloudinit.util import (
-    ensure_dir, load_file, write_file, yaml_dumps)
+    ensure_dir, load_file, write_file)
 from cloudinit.tests.helpers import (
     FilesystemMockingTestCase, wrap_and_call)
 
@@ -39,7 +40,7 @@ class TestMain(FilesystemMockingTestCase):
             ],
             'cloud_init_modules': ['write-files', 'runcmd'],
         }
-        cloud_cfg = yaml_dumps(self.cfg)
+        cloud_cfg = safeyaml.dumps(self.cfg)
         ensure_dir(os.path.join(self.new_root, 'etc', 'cloud'))
         self.cloud_cfg_file = os.path.join(
             self.new_root, 'etc', 'cloud', 'cloud.cfg')
@@ -113,7 +114,7 @@ class TestMain(FilesystemMockingTestCase):
         """When local-hostname metadata is present, call cc_set_hostname."""
         self.cfg['datasource'] = {
             'None': {'metadata': {'local-hostname': 'md-hostname'}}}
-        cloud_cfg = yaml_dumps(self.cfg)
+        cloud_cfg = safeyaml.dumps(self.cfg)
         write_file(self.cloud_cfg_file, cloud_cfg)
         cmdargs = myargs(
             debug=False, files=None, force=False, local=False, reporter=None,