1# coding=utf-8
2
3
4class ParallelTreeManager(object):
5    """
6    Manages output tree by building it from flat test names.
7    """
8
9    def __init__(self, offset):
10        super(ParallelTreeManager, self).__init__()
11        self._max_node_id = offset
12        self.offset = offset
13        self._branches = dict()  # key is test name as tuple, value is tuple of test_id, parent_id
14
15    def _next_node_id(self):
16        self._max_node_id += 1
17        return self._max_node_id
18
19    def level_opened(self, test_as_list, func_to_open):
20        """
21        To be called on test start.
22
23        :param test_as_list: test name splitted as list
24        :param func_to_open: func to be called if test can open new level
25        :return: None if new level opened, or tuple of command client should execute and try opening level again
26         Command is "open" (open provided level) or "close" (close it). Second item is test name as list
27        """
28
29        if tuple(test_as_list) in self._branches:
30            # We have parent, ok
31            func_to_open()
32            return None
33        elif len(test_as_list) == 1:
34            self._branches[tuple(test_as_list)] = (self._next_node_id(), self.offset)
35            func_to_open()
36            return None
37
38        commands = []
39
40        parent_id = self.offset
41        for i in range(len(test_as_list)):
42            tmp_parent_as_list = test_as_list[0:i + 1]
43            try:
44                parent_id, _ = self._branches[tuple(tmp_parent_as_list)]
45            except KeyError:
46                node_id = self._next_node_id()
47                self._branches[tuple(tmp_parent_as_list)] = (node_id, parent_id)
48                parent_id = node_id
49                if tmp_parent_as_list != test_as_list: # Different test opened
50                    commands.append(("open", tmp_parent_as_list))
51        if commands:
52            return commands
53        else:
54            func_to_open()
55
56    def level_closed(self, test_as_list, func_to_close):
57        """
58        To be called on test end or failure.
59
60        See level_opened doc.
61        """
62        func_to_close()
63
64    # Part of contract
65    def get_node_ids(self, test_name):
66        """
67
68        :return: (current_node_id, parent_node_id) or None, None if message must be ignored
69        """
70        try:
71            return self._branches[tuple(test_name.split("."))]
72        except KeyError:
73            return None, None
74