1import unittest
2
3from .pathtree import PathTree
4
5class PathTreeTest(unittest.TestCase):
6    def test_set_depth_1(self):
7        pt = PathTree()
8        pt['/a'] = 1
9
10    def test_set_depth_2(self):
11        pt = PathTree()
12        pt['/a/b'] = 2
13
14    def test_get_no_key(self):
15        pt = PathTree()
16        with self.assertRaises(KeyError):
17            pt['/foo']
18
19    def test_get_depth_1(self):
20        pt = PathTree()
21        pt['/a'] = 1
22        self.assertEqual(1, pt['/a'])
23
24    def test_get_depth_2(self):
25        pt = PathTree()
26        pt['/a/b'] = 2
27        self.assertEqual(set(['/a', '/a/b']), set(pt.keys()))
28        self.assertEqual(2, pt['/a/b'])
29
30    def test_get_default(self):
31        self.assertEquals(1, PathTree().get('/a', 1))
32
33    def test_get_present(self):
34        pt = PathTree()
35        pt['/a'] = 1
36        self.assertEquals(1, pt.get('/a'))
37
38    def test_set_2_depth_1(self):
39        pt = PathTree()
40        pt['/a'] = 1.1
41        pt['/b'] = 1.2
42        self.assertEqual(set(['/a', '/b']), set(pt.keys()))
43        self.assertEqual(1.1, pt['/a'])
44        self.assertEqual(1.2, pt['/b'])
45
46    def test_set_2_depth_2_common_parent(self):
47        pt = PathTree()
48        pt['/a/b'] = 1.1
49        pt['/a/c'] = 1.2
50        self.assertEqual(set(['/a', '/a/b', '/a/c']), set(pt.keys()))
51        self.assertEqual(1.1, pt['/a/b'])
52        self.assertEqual(1.2, pt['/a/c'])
53
54    def test_set_2_depth_2_separate_parent(self):
55        pt = PathTree()
56        pt['/a/b'] = 1.1
57        pt['/b/c'] = 1.2
58        self.assertEqual(set(['/a', '/b', '/a/b', '/b/c']), set(pt.keys()))
59        self.assertEqual(1.1, pt['/a/b'])
60        self.assertEqual(1.2, pt['/b/c'])
61
62    def test_dumpd_empty(self):
63        pt = PathTree()
64        self.assertEquals(dict(), pt.dumpd())
65
66    def test_dumpd_populated(self):
67        pt = PathTree()
68        pt['/a/b'] = { 1 : 1.1 }
69        pt['/b/c'] = { 2 : 1.2 }
70        dump = pt.dumpd()
71        self.assertEquals(set(['/']), set(dump.keys()))
72        self.assertEquals(set(['a', 'b']), set(dump['/'].keys()))
73        self.assertEquals(set(['b']), set(dump['/']['a'].keys()))
74        self.assertEquals(set(['c']), set(dump['/']['b'].keys()))
75
76    def test_del_set_1_depth_1(self):
77        pt = PathTree()
78        pt['/a'] = 1
79        del pt['/a']
80        self.assertEquals(0, len(pt.keys()))
81
82    def test_del_set_2_depth_1(self):
83        pt = PathTree()
84        pt['/a'] = 1
85        pt['/b'] = 2
86        del pt['/a']
87        self.assertEquals(set(['/b']), set(pt.keys()))
88
89    def test_del_set_2_depth_2(self):
90        pt = PathTree()
91        pt['/a/b'] = 1
92        pt['/b/c'] = 2
93        del pt['/a/b']
94        self.assertEquals(set(['/b/c', '/b']), set(pt.keys()))
95
96    def test_setdefault_present(self):
97        pt = PathTree()
98        pt['/a'] = 1
99        self.assertEquals(1, pt.setdefault('/a', 2))
100
101    def test_setdefault_absent(self):
102        self.assertEquals(1, PathTree().setdefault('/a', 1))
103
104    def test_del_no_key(self):
105        with self.assertRaises(KeyError):
106            del PathTree()['/a']
107
108    def test_values_1(self):
109        pt = PathTree()
110        pt['/a'] = 1
111        self.assertEquals(set([1]), set(pt.values()))
112
113    def test_values_2(self):
114        pt = PathTree()
115        pt['/a'] = 1
116        pt['/b'] = 2
117        self.assertEquals(set([1, 2]), set(pt.values()))
118
119    def test_items_1(self):
120        pt = PathTree()
121        pt['/a'] = 1
122        self.assertEquals(set([('/a', 1)]), set(pt.items()))
123
124    def test_items_2(self):
125        pt = PathTree()
126        pt['/a'] = 1
127        pt['/b'] = 2
128        self.assertEquals(set([('/a', 1), ('/b', 2)]), set(pt.items()))
129
130    def test_items_depth_2(self):
131        pt = PathTree()
132        pt['/a/b'] = 1
133        self.assertEquals(set([('/a', None), ('/a/b', 1)]), set(pt.items()))
134
135    def test_dataitems_0(self):
136        pt = PathTree()
137        self.assertEquals(set(), set(pt.dataitems()))
138
139    def test_dataitems_1(self):
140        pt = PathTree()
141        pt['/b'] = 1
142        self.assertEquals(set([('/b', 1)]), set(pt.dataitems()))
143
144    def test_dataitems_2(self):
145        pt = PathTree()
146        pt['/a'] = 1
147        pt['/b/c'] = 2
148        self.assertEquals(set([('/a', 1), ('/b/c', 2)]), set(pt.dataitems()))
149
150    def test_get_children(self):
151        pt = PathTree()
152        pt['/a'] = 1
153        self.assertEquals(set(['a']), set(pt.get_children('/')))
154        self.assertEquals(set(), set(pt.get_children('/a')))
155
156    def test_get_children_nested(self):
157        pt = PathTree()
158        pt['/a/b'] = 1
159        self.assertEquals(set(['a']), set(pt.get_children('/')))
160        self.assertEquals(set(['b']), set(pt.get_children('/a')))
161        self.assertEquals(set(), set(pt.get_children('/a/b')))
162
163    def test_demote_1(self):
164        pt = PathTree()
165        pt['/a'] = 1
166        self.assertEquals([1], pt.values())
167        pt.demote('/a')
168        self.assertEquals([None], pt.values())
169
170    def test_demote_2(self):
171        pt = PathTree()
172        pt['/a'] = 1
173        pt['/b'] = 2
174        self.assertEquals(set([1, 2]), set(pt.values()))
175        pt.demote('/a')
176        self.assertEquals(set([None, 2]), set(pt.values()))
177
178    def test_demote_nested(self):
179        pt = PathTree()
180        pt['/a'] = 1
181        pt['/a/b'] = 2
182        self.assertEquals(set([1, 2]), set(pt.values()))
183        pt.demote('/a/b')
184        self.assertEquals(set([1, None]), set(pt.values()))
185
186    def test_iter(self):
187        pt = PathTree()
188        pt['/a'] = 1
189        i = iter(pt)
190        k, v = next(i)
191        self.assertEquals('/a', k)
192        self.assertEquals(1, v)
193        with self.assertRaises(StopIteration):
194            next(i)
195
196    def test_iter_2(self):
197        pt = PathTree()
198        pt['/a'] = 1
199        pt['/b'] = 2
200        i = iter(pt)
201        k, v = next(i)
202        self.assertEquals('/a', k)
203        self.assertEquals(1, v)
204        k, v = next(i)
205        self.assertEquals('/b', k)
206        self.assertEquals(2, v)
207        with self.assertRaises(StopIteration):
208            next(i)
209
210    def test_iter_2_nested(self):
211        pt = PathTree()
212        pt['/a'] = 1
213        pt['/a/b'] = 2
214        i = iter(pt)
215        k, v = next(i)
216        self.assertEquals('/a', k)
217        self.assertEquals(1, v)
218        k, v = next(i)
219        self.assertEquals('/a/b', k)
220        self.assertEquals(2, v)
221        with self.assertRaises(StopIteration):
222            next(i)
223
224    def test_keys_2_nested_depth_1(self):
225        pt = PathTree()
226        pt['/a'] = 1
227        pt['/a/b'] = 2
228        self.assertEquals(set(['/a']), set(pt.keys(depth=1)))
229
230    def test_values_2_nested_depth_1(self):
231        pt = PathTree()
232        pt['/a'] = 1
233        pt['/a/b'] = 2
234        self.assertEquals(set([1]), set(pt.values(depth=1)))
235
236    def test_items_2_nested_depth_1(self):
237        pt = PathTree()
238        pt['/a'] = 1
239        pt['/a/b'] = 2
240        self.assertEquals(set([('/a', 1)]), set(pt.items(depth=1)))
241
242    def test_dataitems_2_nested_depth_1(self):
243        pt = PathTree()
244        pt['/a'] = 1
245        pt['/a/b'] = 2
246        pt['/b'] = None
247        pt['/b/c'] = 1
248        self.assertEquals(set([('/a', 1)]), set(pt.dataitems(depth=1)))
249
250    def test_keys_2_nested_subtree(self):
251        pt = PathTree()
252        pt['/a'] = 1
253        pt['/a/b'] = 2
254        pt['/b'] = 3
255        self.assertEquals(set(['/a/b']), set(pt.keys(subtree='/a')))
256
257    def test_values_2_nested_subtree(self):
258        pt = PathTree()
259        pt['/a'] = 1
260        pt['/a/b'] = 2
261        pt['/b'] = 3
262        self.assertEquals(set([2]), set(pt.values(subtree='/a')))
263
264    def test_items_2_nested_subtree(self):
265        pt = PathTree()
266        pt['/a'] = 1
267        pt['/a/b'] = 2
268        pt['/b'] = 3
269        self.assertEquals(set([('/a/b', 2)]), set(pt.items(subtree='/a')))
270
271    def test_dataitems_2_nested_subtree(self):
272        pt = PathTree()
273        pt['/a'] = 1
274        pt['/a/b'] = 2
275        pt['/a/c'] = None
276        pt['/b'] = 3
277        self.assertEquals(set([('/a/b', 2)]), set(pt.dataitems(subtree='/a')))
278
279    def test_keys_3_nested_subtree_depth_1(self):
280        pt = PathTree()
281        pt['/a'] = 1
282        pt['/a/b'] = 2
283        pt['/a/b/c'] = 3
284        pt['/b'] = 4
285        self.assertEquals(set(['/a/b']), set(pt.keys(subtree='/a', depth=1)))
286
287    def test_values_3_nested_subtree_depth_1(self):
288        pt = PathTree()
289        pt['/a'] = 1
290        pt['/a/b'] = 2
291        pt['/a/b/c'] = 3
292        pt['/b'] = 4
293        self.assertEquals(set([2]), set(pt.values(subtree='/a', depth=1)))
294
295    def test_items_3_nested_subtree_depth_1(self):
296        pt = PathTree()
297        pt['/a'] = 1
298        pt['/a/b'] = 2
299        pt['/a/b/c'] = 3
300        pt['/b'] = 4
301        self.assertEquals(set([('/a/b', 2)]), set(pt.items(subtree='/a', depth=1)))
302
303    def test_items_3_nested_subtree_depth_1(self):
304        pt = PathTree()
305        pt['/a'] = 1
306        pt['/a/b'] = 2
307        pt['/a/b/c'] = 3
308        pt['/a/c'] = None
309        pt['/b'] = 4
310        self.assertEquals(set([('/a/b', 2)]), set(pt.dataitems(subtree='/a', depth=1)))
311
312import timeit
313import sys
314
315def depth_stress(pt, k):
316    pt[k] = 1
317    pt[k] = pt[k] + 1
318    del pt[k]
319
320depth_setup = """\
321from __main__ import depth_stress
322from obmc.utils.pathtree import PathTree
323pt = PathTree()
324key = '/' + '/'.join(['a'] * {})
325"""
326
327def width_stress(pt, k):
328    pt.get_children(k)
329
330width_setup = """\
331from __main__ import width_stress
332from obmc.utils.pathtree import PathTree
333pt = PathTree()
334for i in range(0, {}):
335    pt["/{}/a".format(i)] = i
336key = '/0/a'
337"""
338
339def iter_stress(pt):
340    for i in pt.dataitems():
341        pass
342
343iter_setup = """\
344from __main__ import iter_stress
345from obmc.utils.pathtree import PathTree
346pt = PathTree()
347for i in ('/' + '/'.join([chr(ord('a') + x)] * {}) for x in range(0, 26)):
348    for j in range(0, {}):
349        k = "{}".format(i, j)
350        pt[k] = k
351"""
352
353if __name__ == "__main__":
354    print("Depth tests:")
355    for depth in range(1, 11):
356        setup = depth_setup.format(depth)
357        stmt = "depth_stress(pt, key)"
358        time = timeit.timeit(stmt, setup=setup)
359        print("\t{}: {}".format(depth, time))
360    print
361    print("Width tests:")
362    for width in map(lambda x: 2 ** x, range(0, 21)):
363        setup = width_setup.format(width, "{}")
364        stmt = "width_stress(pt, key)"
365        time = timeit.timeit(stmt, setup=setup)
366        print("\t{}: {}".format(width, time))
367    print
368    print("Iteration tests:")
369    n = 1000
370    for depth in range(1, 5):
371        for width in range(1, 5):
372            setup = iter_setup.format(depth, width, "{}/{}")
373            stmt = "iter_stress(pt)"
374            time = timeit.timeit(stmt, setup=setup, number=n)
375            print("\tdepth={}, width={}, n={}: {}".format(depth, width, n, time))
376