MsgTreeTest.py

# ClusterShell test suite
# Written by S. Thiell

"""Unit test for ClusterShell MsgTree Class"""

from operator import itemgetter
import sys
import unittest

from ClusterShell.MsgTree import *


class MsgTreeTest(unittest.TestCase):

    def test_001_basics(self):
        """test MsgTree basics"""
        tree = MsgTree()
        self.assertEqual(len(tree), 0)

        tree.add("key", b"message")
        self.assertEqual(len(tree), 1)

        tree.add("key", b"message2")
        self.assertEqual(len(tree), 1)

        tree.add("key2", b"message3")
        self.assertEqual(len(tree), 2)

    def test_002_elem(self):
        """test MsgTreeElem"""
        elem = MsgTreeElem()
        self.assertEqual(len(elem), 0)
        for s in elem:
            self.fail("found line in empty MsgTreeElem!")

    def test_003_iterators(self):
        """test MsgTree iterators"""
        # build tree...
        tree = MsgTree()
        self.assertEqual(len(tree), 0)
        tree.add(("item1", "key"), b"message0")
        self.assertEqual(len(tree), 1)
        tree.add(("item2", "key"), b"message2")
        self.assertEqual(len(tree), 2)
        tree.add(("item3", "key"), b"message3")
        self.assertEqual(len(tree), 3)
        tree.add(("item4", "key"), b"message3")
        tree.add(("item2", "newkey"), b"message4")
        self.assertEqual(len(tree), 5)
        self.assertEqual(tree._depth(), 1)

        # test standard iterator (over keys)
        cnt = 0
        what = set([ ("item1", "key"), ("item2", "key"), ("item3", "key"), \
                    ("item4", "key"), ("item2", "newkey") ])
        for key in tree:
            cnt += 1
            what.remove(key)
        self.assertEqual(cnt, 5)
        self.assertEqual(len(what), 0)

        # test keys() iterator
        cnt = 0
        for key in tree.keys(): # keep this test for return value check
            cnt += 1
        self.assertEqual(cnt, 5)
        self.assertEqual(len(list(iter(tree.keys()))), 5)

        # test messages() iterator (iterate over different messages)
        cnt = 0
        for msg in tree.messages():
            cnt += 1
            self.assertEqual(len(msg), len(b"message0"))
            self.assertEqual(msg[0][:-1], b"message")
        self.assertEqual(cnt, 4)
        self.assertEqual(len(list(iter(tree.messages()))), 4)

        # test items() iterator (iterate over all key, msg pairs)
        cnt = 0
        for key, msg in tree.items():
            cnt += 1
        self.assertEqual(cnt, 5)
        self.assertEqual(len(list(iter(tree.items()))), 5)

        # test walk() iterator (iterate by msg and give the list of
        # associated keys)
        cnt = 0
        cnt_2 = 0
        for msg, keys in tree.walk():
            cnt += 1
            if len(keys) == 2:
                self.assertEqual(msg, b"message3")
                cnt_2 += 1
        self.assertEqual(cnt, 4)
        self.assertEqual(cnt_2, 1)
        self.assertEqual(len(list(iter(tree.walk()))), 4)

        # test walk() with provided key-filter
        cnt = 0
        for msg, keys in tree.walk(match=lambda s: s[1] == "newkey"):
            cnt += 1
        self.assertEqual(cnt, 1)

        # test walk() with provided key-mapper
        cnt = 0
        cnt_2 = 0
        for msg, keys in tree.walk(mapper=itemgetter(0)):
            cnt += 1
            if len(keys) == 2:
                for k in keys:
                    self.assertEqual(type(k), str)
                cnt_2 += 1
        self.assertEqual(cnt, 4)
        self.assertEqual(cnt_2, 1)

        # test walk with full options: key-filter and key-mapper
        cnt = 0
        for msg, keys in tree.walk(match=lambda k: k[1] == "newkey",
                                       mapper=itemgetter(0)):
            cnt += 1
            self.assertEqual(msg, b"message4")
            self.assertEqual(keys[0], "item2")
        self.assertEqual(cnt, 1)

        cnt = 0
        for msg, keys in tree.walk(match=lambda k: k[1] == "key",
                                       mapper=itemgetter(0)):
            cnt += 1
            self.assertEqual(keys[0][:-1], "item")
        self.assertEqual(cnt, 3) # 3 and not 4 because item3 and item4 are merged

    def test_004_getitem(self):
        """test MsgTree get and __getitem__"""
        # build tree...
        tree = MsgTree()
        tree.add("item1", b"message0")
        self.assertEqual(len(tree), 1)
        tree.add("item2", b"message2")
        tree.add("item3", b"message2")
        tree.add("item4", b"message3")
        tree.add("item2", b"message4")
        tree.add("item3", b"message4")
        self.assertEqual(len(tree), 4)
        self.assertEqual(tree["item1"], b"message0")
        self.assertEqual(tree.get("item1"), b"message0")
        self.assertEqual(tree["item2"], b"message2\nmessage4")
        self.assertEqual(tree.get("item2"), b"message2\nmessage4")
        self.assertEqual(tree.get("item5", b"default_buf"), b"default_buf")
        self.assertEqual(tree._depth(), 2)

    def test_005_remove(self):
        """test MsgTree.remove()"""
        # build tree
        tree = MsgTree()
        self.assertEqual(len(tree), 0)
        tree.add(("w1", "key1"), b"message0")
        self.assertEqual(len(tree), 1)
        tree.add(("w1", "key2"), b"message0")
        self.assertEqual(len(tree), 2)
        tree.add(("w1", "key3"), b"message0")
        self.assertEqual(len(tree), 3)
        tree.add(("w2", "key4"), b"message1")
        self.assertEqual(len(tree), 4)
        tree.remove(lambda k: k[1] == "key2")
        self.assertEqual(len(tree), 3)
        for msg, keys in tree.walk(match=lambda k: k[0] == "w1",
                                   mapper=itemgetter(1)):
            self.assertEqual(msg, b"message0")
            self.assertEqual(len(keys), 2)
        tree.remove(lambda k: k[0] == "w1")
        self.assertEqual(len(tree), 1)
        tree.remove(lambda k: k[0] == "w2")
        self.assertEqual(len(tree), 0)
        tree.clear()
        self.assertEqual(len(tree), 0)

    def test_006_scalability(self):
        """test MsgTree scalability"""
        # test tree of 10k nodes with a single different line each
        tree = MsgTree()
        for i in range(10000):
            tree.add("node%d" % i, b"message" + str(i).encode('ascii'))
        self.assertEqual(len(tree), 10000)
        cnt = 0
        for msg, keys in tree.walk():
            cnt += 1
        self.assertEqual(cnt, 10000)

        # test tree of 1 node with 10k lines
        tree = MsgTree()
        for i in range(10000):
            tree.add("nodeX", b"message" + str(i).encode('ascii'))
        self.assertEqual(len(tree), 1)
        cnt = 0
        for msg, keys in tree.walk():
            testlines = bytes(msg) # test MsgTreeElem.__iter__()
            self.assertEqual(len(testlines.splitlines()), 10000)
            cnt += 1
        self.assertEqual(cnt, 1)

        # test tree of 100 nodes with the same 1000 message lines each
        tree = MsgTree()
        for j in range(100):
            for i in range(1000):
                tree.add("node%d" % j, b"message" + str(i).encode('ascii'))
        self.assertEqual(len(tree), 100)
        cnt = 0
        for msg, keys in tree.walk():
            testlines = bytes(msg) # test MsgTreeElem.__iter__()
            self.assertEqual(len(testlines.splitlines()), 1000)
            cnt += 1
        self.assertEqual(cnt, 1)

    def test_007_shift_mode(self):
        """test MsgTree in shift mode"""
        tree = MsgTree(mode=MODE_SHIFT)
        tree.add("item1", b"message0")
        self.assertEqual(len(tree), 1)
        tree.add("item2", b"message2")
        tree.add("item3", b"message2")
        tree.add("item4", b"message3")
        tree.add("item2", b"message4")
        tree.add("item3", b"message4")
        self.assertEqual(len(tree), 4)
        self.assertEqual(tree["item1"], b"message0")
        self.assertEqual(tree.get("item1"), b"message0")
        self.assertEqual(tree["item2"], b"message2\nmessage4")
        self.assertEqual(tree.get("item2"), b"message2\nmessage4")
        self.assertEqual(tree.get("item5", b"default_buf"), b"default_buf")
        self.assertEqual(tree._depth(), 2)
        self.assertEqual(len(list(tree.walk())), 3)

    def test_008_trace_mode(self):
        """test MsgTree in trace mode"""
        tree = MsgTree(mode=MODE_TRACE)
        tree.add("item1", b"message0")
        self.assertEqual(len(tree), 1)
        tree.add("item2", b"message2")
        tree.add("item3", b"message2")
        tree.add("item4", b"message3")
        tree.add("item2", b"message4")
        tree.add("item3", b"message4")
        self.assertEqual(len(tree), 4)
        self.assertEqual(tree["item1"], b"message0")
        self.assertEqual(tree.get("item1"), b"message0")
        self.assertEqual(tree["item2"], b"message2\nmessage4")
        self.assertEqual(tree.get("item2"), b"message2\nmessage4")
        self.assertEqual(tree.get("item5", b"default_buf"), b"default_buf")
        self.assertEqual(tree._depth(), 2)
        self.assertEqual(len(list(tree.walk())), 4)
        # /!\ results are not sorted
        result = [(m, sorted(k), d, c) for m, k, d, c
                  in sorted(list(tree.walk_trace()))]
        self.assertEqual(result,
                         [(b'message0', ['item1'], 1, 0),
                          (b'message2', ['item2', 'item3'], 1, 1),
                          (b'message3', ['item4'], 1, 0),
                          (b'message4', ['item2', 'item3'], 2, 0)])

    def test_009_defer_to_shift_mode(self):
        """test MsgTree defer to shift mode"""
        tree = MsgTree(mode=MODE_DEFER)
        tree.add("item1", b"message0")
        self.assertEqual(len(tree), 1)
        tree.add("item2", b"message1")
        self.assertEqual(len(tree), 2)
        tree.add("item3", b"message2")
        self.assertEqual(len(tree), 3)
        tree.add("item2", b"message4")
        tree.add("item1", b"message3")
        self.assertEqual(tree["item1"], b"message0\nmessage3")
        self.assertEqual(tree.mode, MODE_DEFER)
        # calling walk with call _update_keys() and change to MODE_SHIFT
        self.assertEqual(sorted([(k, e.message()) for e, k in tree.walk()]),
                         [(['item1'], b'message0\nmessage3'),
                          (['item2'], b'message1\nmessage4'),
                          (['item3'], b'message2')])
        self.assertEqual(tree.mode, MODE_SHIFT)
        # further tree modifications should be safe...
        tree.add("item1", b"message5")
        tree.add("item2", b"message6")
        self.assertEqual(tree["item1"], b"message0\nmessage3\nmessage5")
        # /!\ results are not sorted
        self.assertEqual(sorted([(k, e.message()) for e, k in tree.walk()]),
                         [(['item1'], b'message0\nmessage3\nmessage5'),
                          (['item2'], b'message1\nmessage4\nmessage6'),
                          (['item3'], b'message2')])

    def test_010_remove_in_defer_mode(self):
        """test MsgTree remove in defer mode"""
        tree = MsgTree(mode=MODE_DEFER)
        tree.add("item1", b"message0")
        self.assertEqual(len(tree), 1)
        tree.add("item2", b"message1")
        self.assertEqual(len(tree), 2)
        tree.add("item3", b"message2")
        self.assertEqual(len(tree), 3)
        tree.add("item2", b"message4")
        tree.add("item1", b"message3")
        tree.remove(lambda k: k == "item2")
        self.assertEqual(tree["item1"], b"message0\nmessage3")
        self.assertRaises(KeyError, tree.__getitem__, "item2")
        # calling walk with call _update_keys() and change to MODE_SHIFT
        self.assertEqual(sorted([(k, e.message()) for e, k in tree.walk()]),
                         [(['item1'], b'message0\nmessage3'),
                          (['item3'], b'message2')])
        self.assertEqual(tree.mode, MODE_SHIFT)
        # further tree modifications should be safe...
        tree.add("item1", b"message5")
        tree.add("item2", b"message6")
        self.assertEqual(tree["item1"], b"message0\nmessage3\nmessage5")
        self.assertEqual(tree["item2"], b"message6")
        # /!\ results are not sorted
        self.assertEqual(sorted([(k, e.message()) for e, k in tree.walk()]),
                         [(['item1'], b'message0\nmessage3\nmessage5'),
                          (['item2'], b'message6'),
                          (['item3'], b'message2')])

    def test_011_str_compat(self):
        """test MsgTreeElem.__str__() compatibility"""
        tree = MsgTree()
        tree.add("item1", b"message0")
        elem = tree["item1"]
        if sys.version_info >= (3, 0):
            # casting to string is definitively not supported in Python 3,
            # use bytes() instead.
            self.assertRaises(TypeError, elem.__str__)
            self.assertEqual(bytes(elem), b"message0")
        else:
            self.assertEqual(str(elem), "message0")