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