1from __future__ import absolute_import, division, print_function 2 3__metaclass__ = type 4 5import io 6import yaml 7 8from ansible.module_utils.six import PY3 9from ansible.parsing.yaml.loader import AnsibleLoader 10from ansible.parsing.yaml.dumper import AnsibleDumper 11 12 13class YamlTestUtils(object): 14 """Mixin class to combine with a unittest.TestCase subclass.""" 15 16 def _loader(self, stream): 17 """Vault related tests will want to override this. 18 19 Vault cases should setup a AnsibleLoader that has the vault password.""" 20 return AnsibleLoader(stream) 21 22 def _dump_stream(self, obj, stream, dumper=None): 23 """Dump to a py2-unicode or py3-string stream.""" 24 if PY3: 25 return yaml.dump(obj, stream, Dumper=dumper) 26 else: 27 return yaml.dump(obj, stream, Dumper=dumper, encoding=None) 28 29 def _dump_string(self, obj, dumper=None): 30 """Dump to a py2-unicode or py3-string""" 31 if PY3: 32 return yaml.dump(obj, Dumper=dumper) 33 else: 34 return yaml.dump(obj, Dumper=dumper, encoding=None) 35 36 def _dump_load_cycle(self, obj): 37 # Each pass though a dump or load revs the 'generation' 38 # obj to yaml string 39 string_from_object_dump = self._dump_string(obj, dumper=AnsibleDumper) 40 41 # wrap a stream/file like StringIO around that yaml 42 stream_from_object_dump = io.StringIO(string_from_object_dump) 43 loader = self._loader(stream_from_object_dump) 44 # load the yaml stream to create a new instance of the object (gen 2) 45 obj_2 = loader.get_data() 46 47 # dump the gen 2 objects directory to strings 48 string_from_object_dump_2 = self._dump_string( 49 obj_2, dumper=AnsibleDumper 50 ) 51 52 # The gen 1 and gen 2 yaml strings 53 self.assertEqual(string_from_object_dump, string_from_object_dump_2) 54 # the gen 1 (orig) and gen 2 py object 55 self.assertEqual(obj, obj_2) 56 57 # again! gen 3... load strings into py objects 58 stream_3 = io.StringIO(string_from_object_dump_2) 59 loader_3 = self._loader(stream_3) 60 obj_3 = loader_3.get_data() 61 62 string_from_object_dump_3 = self._dump_string( 63 obj_3, dumper=AnsibleDumper 64 ) 65 66 self.assertEqual(obj, obj_3) 67 # should be transitive, but... 68 self.assertEqual(obj_2, obj_3) 69 self.assertEqual(string_from_object_dump, string_from_object_dump_3) 70 71 def _old_dump_load_cycle(self, obj): 72 """Dump the passed in object to yaml, load it back up, dump again, compare.""" 73 stream = io.StringIO() 74 75 yaml_string = self._dump_string(obj, dumper=AnsibleDumper) 76 self._dump_stream(obj, stream, dumper=AnsibleDumper) 77 78 yaml_string_from_stream = stream.getvalue() 79 80 # reset stream 81 stream.seek(0) 82 83 loader = self._loader(stream) 84 # loader = AnsibleLoader(stream, vault_password=self.vault_password) 85 obj_from_stream = loader.get_data() 86 87 stream_from_string = io.StringIO(yaml_string) 88 loader2 = self._loader(stream_from_string) 89 # loader2 = AnsibleLoader(stream_from_string, vault_password=self.vault_password) 90 obj_from_string = loader2.get_data() 91 92 stream_obj_from_stream = io.StringIO() 93 stream_obj_from_string = io.StringIO() 94 95 if PY3: 96 yaml.dump( 97 obj_from_stream, stream_obj_from_stream, Dumper=AnsibleDumper 98 ) 99 yaml.dump( 100 obj_from_stream, stream_obj_from_string, Dumper=AnsibleDumper 101 ) 102 else: 103 yaml.dump( 104 obj_from_stream, 105 stream_obj_from_stream, 106 Dumper=AnsibleDumper, 107 encoding=None, 108 ) 109 yaml.dump( 110 obj_from_stream, 111 stream_obj_from_string, 112 Dumper=AnsibleDumper, 113 encoding=None, 114 ) 115 116 yaml_string_stream_obj_from_stream = stream_obj_from_stream.getvalue() 117 yaml_string_stream_obj_from_string = stream_obj_from_string.getvalue() 118 119 stream_obj_from_stream.seek(0) 120 stream_obj_from_string.seek(0) 121 122 if PY3: 123 yaml_string_obj_from_stream = yaml.dump( 124 obj_from_stream, Dumper=AnsibleDumper 125 ) 126 yaml_string_obj_from_string = yaml.dump( 127 obj_from_string, Dumper=AnsibleDumper 128 ) 129 else: 130 yaml_string_obj_from_stream = yaml.dump( 131 obj_from_stream, Dumper=AnsibleDumper, encoding=None 132 ) 133 yaml_string_obj_from_string = yaml.dump( 134 obj_from_string, Dumper=AnsibleDumper, encoding=None 135 ) 136 137 assert yaml_string == yaml_string_obj_from_stream 138 assert ( 139 yaml_string 140 == yaml_string_obj_from_stream 141 == yaml_string_obj_from_string 142 ) 143 assert ( 144 yaml_string 145 == yaml_string_obj_from_stream 146 == yaml_string_obj_from_string 147 == yaml_string_stream_obj_from_stream 148 == yaml_string_stream_obj_from_string 149 ) 150 assert obj == obj_from_stream 151 assert obj == obj_from_string 152 assert obj == yaml_string_obj_from_stream 153 assert obj == yaml_string_obj_from_string 154 assert ( 155 obj 156 == obj_from_stream 157 == obj_from_string 158 == yaml_string_obj_from_stream 159 == yaml_string_obj_from_string 160 ) 161 return { 162 "obj": obj, 163 "yaml_string": yaml_string, 164 "yaml_string_from_stream": yaml_string_from_stream, 165 "obj_from_stream": obj_from_stream, 166 "obj_from_string": obj_from_string, 167 "yaml_string_obj_from_string": yaml_string_obj_from_string, 168 } 169