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