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