unittests.py

# -*- coding: utf-8 -*-
"""
    This is part of pyahocorasick Python module.

    Unit tests for the C-based ahocorasick module.

    Author    : Wojciech Muła, wojciech_mula@poczta.onet.pl
    WWW       : http://0x80.pl/proj/pyahocorasick/
    License   : public domain
"""

import sys
import os
import unittest
import ahocorasick

try:
    import _pickle
except ImportError:
    _pickle = None


if ahocorasick.unicode:
    conv = lambda x: x
else:
    if sys.version_info.major >= 3:
        conv = lambda x: bytes(x, 'ascii')
    else:
        conv = lambda x: x


class TestCase(unittest.TestCase):
    def __init__(self, *args):
        super(TestCase, self).__init__(*args)

        if not hasattr(self, 'assertRaisesRegex'):
            # fixup for Py2
            self.assertRaisesRegex = self.assertRaisesRegexp


    def assertEmpty(self, collection):
        self.assertEqual(0, len(collection))


    def assertNotEmpty(self, collection):
        self.assertGreater(len(collection), 0)


class TestConstructor(TestCase):
    def test_constructor_wrong_store(self):
        with self.assertRaisesRegex(ValueError, "store value must be one of.*"):
            ahocorasick.Automaton(-42)


    def test_constructor_wrong_key_type(self):
        with self.assertRaisesRegex(ValueError, "key_type must have value.*"):
            ahocorasick.Automaton(ahocorasick.STORE_ANY, -42)


class TestTrieStorePyObjectsBase(TestCase):
    def setUp(self):
        self.A = ahocorasick.Automaton();
        self.words = "word python aho corasick \x00\x00\x00".split()
        self.inexisting = "test foo bar dword".split()


class TestTrieMethods(TestTrieStorePyObjectsBase):
    "Test basic methods related to trie structure"

    def test_empty(self):
        A = self.A
        self.assertTrue(A.kind == ahocorasick.EMPTY)
        self.assertTrue(len(A) == 0)


    def test_add_word(self):
        A = self.A
        self.assertTrue(A.kind == ahocorasick.EMPTY)

        n = 0
        for word in self.words:
            n += 1
            A.add_word(conv(word), None)
            self.assertEqual(A.kind, ahocorasick.TRIE)
            self.assertEqual(len(A), n)

        # dupliacted entry
        A.add_word(conv(self.words[0]), None)
        self.assertTrue(A.kind == ahocorasick.TRIE)
        self.assertTrue(len(A) == n)


    def test_add_empty_word(self):
        if ahocorasick.unicode:
            self.assertFalse(self.A.add_word("", None))
        else:
            self.assertFalse(self.A.add_word(b"", None))

        self.assertEqual(len(self.A), 0)
        self.assertEqual(self.A.kind, ahocorasick.EMPTY)


    def test_clear(self):
        A = self.A
        self.assertTrue(A.kind == ahocorasick.EMPTY)

        for w in self.words:
            A.add_word(conv(w), w)

        self.assertEqual(len(A), len(self.words))

        A.clear()
        self.assertEqual(A.kind, ahocorasick.EMPTY)
        self.assertEqual(len(A), 0)


    def test_exists(self):
        A = self.A

        for w in self.words:
            A.add_word(conv(w), w)

        for w in self.words:
            self.assertTrue(A.exists(conv(w)))

        for w in self.inexisting:
            self.assertFalse(A.exists(conv(w)))


    def test_contains(self):
        A = self.A
        for w in self.words:
            A.add_word(conv(w), w)

        for w in self.words:
            self.assertTrue(conv(w) in A)

        for w in self.inexisting:
            self.assertTrue(conv(w) not in A)


    def test_match(self):
        A = self.A
        for word in self.words:
            A.add_word(conv(word), word)

        prefixes = "w wo wor word p py pyt pyth pytho python \x00 \x00\x00 \x00\x00\x00".split()

        for word in prefixes:
            self.assertTrue(A.match(conv(word)))

        inexisting = "wa apple pyTon \x00\x00\x00\x00".split()
        for word in inexisting:
            self.assertFalse(A.match(conv(word)))


    def test_get1(self):
        A = self.A
        for i, w in enumerate(self.words):
            A.add_word(conv(w), i + 1)

        for i, w in enumerate(self.words):
            self.assertEqual(A.get(conv(w)), i + 1)


    def test_get2(self):
        A = self.A
        for i, w in enumerate(self.words):
            A.add_word(conv(w), i + 1)

        for w in self.inexisting:
            self.assertEqual(A.get(conv(w), None), None)


    def test_get3(self):
        A = self.A
        for i, w in enumerate(self.words):
            A.add_word(conv(w), i + 1)

        for w in self.inexisting:
            with self.assertRaises(KeyError):
                A.get(conv(w))


    def test_get_from_an_empty_automaton(self):
        A = ahocorasick.Automaton()

        r = A.get('foo', None)
        self.assertEqual(r, None)


    def test_longest_prefix(self):
        A = self.A
        for i, w in enumerate(self.words):
            A.add_word(conv(w), i + 1)

        # there is "word"
        self.assertEqual(A.longest_prefix(conv("wo")), 2)
        self.assertEqual(A.longest_prefix(conv("working")), 3)
        self.assertEqual(A.longest_prefix(conv("word")), 4)
        self.assertEqual(A.longest_prefix(conv("wordbook")), 4)
        self.assertEqual(A.longest_prefix(conv("void")), 0)
        self.assertEqual(A.longest_prefix(conv("")), 0)


    def test_stats_have_valid_structure(self):
        A = self.A
        for i, w in enumerate(self.words):
            A.add_word(conv(w), i + 1)

        platform_dependent = None
        reference = {
            'longest_word': 8,
            'total_size': platform_dependent,
            'sizeof_node': platform_dependent,
            'nodes_count': 25,
            'words_count': 5,
            'links_count': 24
        }

        s = A.get_stats()

        self.assertEqual(len(s), len(reference))

        for key in reference:
            self.assertIn(key, s)

        for key in (key for key in reference if reference[key] != platform_dependent):
            self.assertEqual(reference[key], s[key])


    def test_stats_for_empty_tire_are_empty(self):
        s = self.A.get_stats()
        self.assertTrue(len(s) > 0)
        for key in s:
            if key != "sizeof_node":
                self.assertEqual(s[key], 0)


class TestTrieRemoveWord(TestTrieStorePyObjectsBase):

    def test_remove_word_from_empty_trie(self):
        self.assertFalse(self.A.remove_word("test"))


    def test_remove_existing_word(self):
        A = self.A

        words = ["he", "her", "hi", "him", "his"]
        for w in words:
            A.add_word(conv(w), w)

        expected_len = len(A)
        for w in words:
            self.assertTrue(self.A.remove_word(w))
            self.assertFalse(self.A.exists(w))
            expected_len -= 1
            self.assertEqual(expected_len, len(A))


    def test_remove_inexisting_word(self):
        A = self.A

        words = ["he", "her", "hi", "him", "his"]
        for w in words:
            A.add_word(conv(w), w)

        expected_len = len(A)
        for w in ["cat", "dog", "tree"]:
            self.assertFalse(self.A.exists(w))
            self.assertFalse(self.A.remove_word(w))
            self.assertEqual(expected_len, len(A))


    def test_remove__case1(self):
        words = ["k", "ki", "kit", "kitt", "kitte", "kitten"
                                 , "kitc", "kitch", "kitche", "kitchen"]

        A = self.A
        for w in words:
            A.add_word(conv(w), w)

        expected_set = set(words)
        for w in words:
            self.assertTrue(self.A.remove_word(w))
            expected_set.discard(w)
            current_set = set(A.keys())
            self.assertEqual(expected_set, current_set)
            self.assertEqual(len(expected_set), len(A))


    def test_remove__case2(self):
        words = ["k", "ki", "kit", "kitt", "kitte", "kitten"
                                 , "kitc", "kitch", "kitche", "kitchen"]

        A = self.A
        for w in words:
            A.add_word(conv(w), w)

        expected_set = set(words)
        for w in reversed(words):
            self.assertTrue(self.A.remove_word(w))
            expected_set.discard(w)
            current_set = set(A.keys())
            self.assertEqual(expected_set, current_set)
            self.assertEqual(len(expected_set), len(A))


    def test_remove_word_changes_type_of_automaton(self):
        A = self.A

        words = ["he", "her", "hi", "him", "his"]
        for w in words:
            A.add_word(conv(w), w)

        A.make_automaton()
        self.assertEqual(ahocorasick.AHOCORASICK, A.kind)

        self.assertFalse(A.remove_word("inexisting"))
        self.assertEqual(ahocorasick.AHOCORASICK, A.kind)

        self.assertTrue(A.remove_word("hi"))
        self.assertEqual(ahocorasick.TRIE, A.kind)


class TestTriePop(TestTrieStorePyObjectsBase):

    def test_pop_from_empty_trie(self):
        with self.assertRaises(KeyError):
            self.A.pop("test")


    def test_pop_existing_word(self):
        A = self.A

        words = ["he", "her", "hi", "him", "his"]
        for w in words:
            A.add_word(conv(w), w)

        expected_len = len(A)
        for w in words:
            self.assertEqual(w, self.A.pop(w))
            self.assertFalse(self.A.exists(w))
            expected_len -= 1
            self.assertEqual(expected_len, len(A))


    def test_pop_inexisting_word(self):
        A = self.A

        words = ["he", "her", "hi", "him", "his"]
        for w in words:
            A.add_word(conv(w), w)

        expected_len = len(A)
        for w in ["cat", "dog", "tree"]:
            with self.assertRaises(KeyError):
                self.A.pop(w)

            self.assertEqual(expected_len, len(A))


    def test_pop__case1(self):
        words = ["k", "ki", "kit", "kitt", "kitte", "kitten"
                                 , "kitc", "kitch", "kitche", "kitchen"]

        A = self.A
        for w in words:
            A.add_word(conv(w), w)

        expected_set = set(words)
        for w in words:
            self.assertEqual(w, self.A.pop(w))
            expected_set.discard(w)
            current_set = set(A.keys())
            self.assertEqual(expected_set, current_set)
            self.assertEqual(len(expected_set), len(A))


    def test_pop__case2(self):
        words = ["k", "ki", "kit", "kitt", "kitte", "kitten"
                                 , "kitc", "kitch", "kitche", "kitchen"]

        A = self.A
        for w in words:
            A.add_word(conv(w), w)

        expected_set = set(words)
        for w in reversed(words):
            self.assertEqual(w, self.A.pop(w))
            expected_set.discard(w)
            current_set = set(A.keys())
            self.assertEqual(expected_set, current_set)
            self.assertEqual(len(expected_set), len(A))


    def test_pop_changes_type_of_automaton(self):
        A = self.A

        words = ["he", "her", "hi", "him", "his"]
        for w in words:
            A.add_word(conv(w), w)

        A.make_automaton()
        self.assertEqual(ahocorasick.AHOCORASICK, A.kind)

        with self.assertRaises(KeyError):
            A.pop("inexisting")

        self.assertEqual(ahocorasick.AHOCORASICK, A.kind)

        self.assertEqual("hi", A.pop("hi"))
        self.assertEqual(ahocorasick.TRIE, A.kind)


class TestTrieIterators(TestTrieStorePyObjectsBase):
    "Test iterators walking over trie"


    def test_iter(self):
        A = self.A
        for i, w in enumerate(self.words):
            A.add_word(conv(w), i + 1)

        L = [word for word in A]
        K = list(map(conv, self.words))
        self.assertEqual(len(L), len(K))
        self.assertEqual(set(L), set(K))


    def test_keys(self):
        A = self.A
        for i, w in enumerate(self.words):
            A.add_word(conv(w), i + 1)

        L = [word for word in A.keys()]
        K = [conv(word) for word in self.words]
        self.assertEqual(len(L), len(K))
        self.assertEqual(set(L), set(K))


    def test_values(self):
        A = self.A
        for i, w in enumerate(self.words):
            A.add_word(conv(w), i + 1)

        L = [x for x in A.values()]
        V = list(range(1, len(self.words) + 1))
        self.assertEqual(len(L), len(V))
        self.assertEqual(set(L), set(V))


    def test_items(self):
        A = self.A
        I = []
        for i, w in enumerate(self.words):
            A.add_word(conv(w), i + 1)
            I.append((conv(w), i + 1))

        L = [x for x in A.items()]
        self.assertEqual(len(L), len(I))
        self.assertEqual(set(L), set(I))


    def test_items_with_prefix_valid(self):
        A = self.A
        words = "he she her hers star ham".split()
        for word in words:
            A.add_word(conv(word), word)

        I = list(map(conv, "he her hers".split()))
        L = [x for x in A.keys(conv("he"))]
        self.assertEqual(len(L), len(I))
        self.assertEqual(set(L), set(I))


    def test_items_with_prefix_invalid(self):
        A = self.A
        words = "he she her hers star ham".split()
        for word in words:
            A.add_word(conv(word), word)

        I = []
        L = [x for x in A.keys(conv("cat"))]
        self.assertEqual(len(L), len(I))
        self.assertEqual(set(L), set(I))


    def test_items_with_valid_pattern(self):
        A = self.A
        words = "abcde aXcd aZcdef aYc Xbcdefgh".split()
        for word in words:
            A.add_word(conv(word), word)

        I = ["aXcd"]
        L = [x for x in A.keys(conv("a?cd"), conv("?"))]
        self.assertEqual(set(I), set(L))


    def test_items_with_valid_pattern2(self):
        A = self.A
        words = "abcde aXcde aZcdef aYc Xbcdefgh".split()
        for word in words:
            A.add_word(conv(word), word)

        L = [x for x in A.keys(conv("a?c??"), conv("?"), ahocorasick.MATCH_EXACT_LENGTH)]
        I = ["abcde", "aXcde"]
        self.assertEqual(set(I), set(L))

        L = [x for x in A.keys(conv("a?c??"), conv("?"), ahocorasick.MATCH_AT_MOST_PREFIX)]
        I = ["aYc", "abcde", "aXcde"]
        self.assertEqual(set(I), set(L))

        L = [x for x in A.keys(conv("a?c??"), conv("?"), ahocorasick.MATCH_AT_LEAST_PREFIX)]
        I = ["abcde", "aXcde", "aZcdef"]
        self.assertEqual(set(I), set(L))


    def test_items_wrong_wildcrard(self):
        with self.assertRaisesRegex(ValueError, "Wildcard must be a single character.*"):
            self.A.keys(conv("anything"), conv("??"))


    def test_items_wrong_match_enum(self):
        with self.assertRaisesRegex(ValueError, "The optional how third argument must be one of"):
            self.A.keys(conv("anything"), conv("?"), -42)


class TestTrieIteratorsInvalidate(TestTrieStorePyObjectsBase):
    "Test invalidating iterator when trie is changed"

    def helper(self, method):
        A = self.A
        for i, w in enumerate(self.words):
            A.add_word(conv(w), i + 1)

        it = method()
        w = next(it)
        # word already exists, just change associated value
        # iterator is still valid
        A.add_word(conv(self.words[0]), 2)
        w = next(it)

        # new word, iterator is invalidated
        A.add_word(conv("should fail"), 1)
        with self.assertRaises(ValueError):
            w = next(it)


    def test_keys(self):
        self.helper(self.A.keys)


    def test_values(self):
        self.helper(self.A.values)


    def test_items(self):
        self.helper(self.A.items)


class TestAutomatonBase(TestCase):
    def setUp(self):
        self.A = ahocorasick.Automaton();
        self.words = "he her hers she".split()
        self.string = "_sherhershe_"
        self.correct_positons = [
            (3, "she"),
            (3, "he"),
            (4, "her"),
            (6, "he"),
            (7, "her"),
            (8, "hers"),
            (10, "she"),
            (10, "he")
        ]


    def add_words(self):
        for word in self.words:
            self.A.add_word(conv(word), word)

        return self.A


    def add_words_and_make_automaton(self):
        self.add_words()
        self.A.make_automaton()
        return self.A


class TestAutomatonConstruction(TestAutomatonBase):
    "Test converting trie to Aho-Corasick automaton"

    def test_make_automaton1(self):
        A = self.A
        self.assertEqual(A.kind, ahocorasick.EMPTY)
        A.make_automaton()
        # empty trie is never converted to automaton
        self.assertEqual(A.kind, ahocorasick.EMPTY)


    def test_make_automaton2(self):
        A = self.A
        self.assertEqual(A.kind, ahocorasick.EMPTY)

        self.add_words()
        self.assertEqual(A.kind, ahocorasick.TRIE)

        A.make_automaton()
        self.assertEqual(A.kind, ahocorasick.AHOCORASICK)


    def test_make_automaton3(self):
        A = self.A
        self.assertEqual(A.kind, ahocorasick.EMPTY)

        self.add_words()
        self.assertEqual(A.kind, ahocorasick.TRIE)

        A.make_automaton()
        self.assertEqual(A.kind, ahocorasick.AHOCORASICK)

        A.add_word(conv("rollback?"), True)
        self.assertEqual(A.kind, ahocorasick.TRIE)


class TestAutomatonSearch(TestAutomatonBase):
    "Test searching using constructed automaton (method find_all)"

    def test_find_all1(self):
        "no action is performed until automaton is constructed"
        A = self.A
        self.assertEqual(A.kind, ahocorasick.EMPTY)

        self.assertEqual(A.find_all(self.string, conv("any arg")), None)

        A.add_word(conv("word"), None)
        self.assertEqual(A.kind, ahocorasick.TRIE)
        self.assertEqual(A.find_all(self.string, conv("any arg")), None)


    def test_find_all2(self):
        A = self.add_words_and_make_automaton()

        L = []
        def callback(index, word):
            L.append((index, word))

        A.find_all(conv(self.string), callback)

        C = self.correct_positons
        self.assertEqual(L, C)


    def test_find_all3(self):
        A = self.add_words_and_make_automaton()

        L = []
        def callback(index, word):
            L.append((index, word))

        start = 4
        end = 9

        L = []
        A.find_all(conv(self.string[start:end]), callback)
        C = [(pos + start, word) for pos, word in L]

        L = []
        A.find_all(conv(self.string), callback, start, end)

        self.assertEqual(L, C)


    def test_find_all__not_a_callable_object(self):
        A = self.add_words_and_make_automaton()

        with self.assertRaisesRegex(TypeError, "The callback argument must be a callable such as a function."):
            A.find_all(conv(self.string), None)


    def test_find_all__wrong_range__case_1(self):
        A = self.add_words_and_make_automaton()

        L = []
        def callback(index, word):
            L.append((index, word))

        with self.assertRaisesRegex(IndexError, "end index not in range 0..12"):
            A.find_all(conv(self.string), callback, 0, len(self.string) + 5)


    def test_find_all__wrong_range__case_2(self):
        A = self.add_words_and_make_automaton()

        L = []
        def callback(index, word):
            L.append((index, word))

        with self.assertRaisesRegex(IndexError, "start index not in range 0..12"):
            A.find_all(conv(self.string), callback, -len(self.string) - 1, 3)


    def test_find_all__end_index_not_given(self):
        A = self.add_words_and_make_automaton()

        L = []
        def callback(index, word):
            L.append((index, word))

        A.find_all(conv(self.string), callback, 0)


    def test_find_all__start_is_negative(self):
        A = self.add_words_and_make_automaton()

        L = []
        def callback(index, word):
            L.append((index, word))

        A.find_all(conv(self.string), callback, -3, 4)


    def test_find_all__end_is_negative(self):
        A = self.add_words_and_make_automaton()

        L = []
        def callback(index, word):
            L.append((index, word))

        A.find_all(conv(self.string), callback, 0, -1)


class TestAutomatonIterSearch(TestAutomatonBase):
    "Test searching using constructed automaton (iterator)"

    def test_iter1(self):
        A = self.A
        self.assertEqual(A.kind, ahocorasick.EMPTY)
        with self.assertRaises(AttributeError):
            A.iter(conv(self.string))

        A.add_word(conv("word"), None)
        self.assertEqual(A.kind, ahocorasick.TRIE)
        with self.assertRaises(AttributeError):
            A.iter(conv(self.string))


    def test_iter2(self):
        A = self.add_words_and_make_automaton()

        L = []
        for index, word in A.iter(conv(self.string)):
            L.append((index, word))

        C = self.correct_positons
        self.assertEqual(L, C)


    def test_iter3(self):
        A = self.add_words_and_make_automaton()

        start = 4
        end = 9

        C = []
        for index, word in A.iter(conv(self.string[start:end])):
            C.append((index + start, word))

        L = []
        for index, word in A.iter(conv(self.string), start, end):
            L.append((index, word))

        self.assertEqual(L, C)


    def test_iter_set(self):
        A = self.add_words_and_make_automaton()
        parts = "_sh erhe rshe _".split()

        expected = {
            '_sh'   : [],
            'erhe'  : [(3, 'she'),
                       (3, 'he'),
                       (4, 'her'),
                       (6, 'he')],
            'rshe'  : [(7, 'her'),
                       (8, 'hers'),
                       (10, 'she'),
                       (10, 'he')],
             '_'    : []
        }

        it = A.iter(conv(""))
        result = {}
        for part in parts:
            it.set(conv(part))
            result[part] = []
            for item in it:
                result[part].append(item)

        self.assertEqual(expected, result)

    def test_iter_set__with_reset(self):
        A = self.add_words_and_make_automaton()

        expected = {
            'he'    : [(1, 'he')],
            'she'   : [(2, 'she'), (2, 'he')],
        }

        it = A.iter(conv(""))
        result = {}
        for part in ["he", "she"]:
            it.set(conv(part), True)
            result[part] = []
            for item in it:
                result[part].append(item)

        self.assertEqual(expected, result)


    def test_iter_compare_with_find_all(self):
        A = self.add_words_and_make_automaton()

        # results from find_all
        L = []
        def callback(index, word):
            L.append((index, word))

        A.find_all(conv(self.string), callback)

        # results from iterator
        C = []
        for index, word in A.iter(conv(self.string)):
            C.append((index, word))

        self.assertEqual(L, C)


    def test_iter_wrong_argument_type(self):
        A = self.add_words_and_make_automaton()

        with self.assertRaisesRegex(TypeError, "string required"):
            A.iter(None)


class TestAutomatonIterSearchWithIgnoreWhiteSpace(TestAutomatonBase):
    "Test searching using constructed automaton (iterator)"

    def setUp(self):
        self.A = ahocorasick.Automaton()
        self.words = "he her hers she".split()
        self.string = "_sh e rher she_"
        self.correct_positons = [
            (4, "she"),
            (4, "he"),
            (6, "her"),
            (8, "he"),
            (9, "her"),
            (11, "hers"),
            (13, "she"),
            (13, "he")
        ]
        self.correct_positons_start_12 = [
            (13, "he")
        ]


    def test_iter1(self):
        self.add_words_and_make_automaton()
        A = self.A
        self.assertEqual(A.kind, ahocorasick.AHOCORASICK)

        L = []
        for index, word in A.iter(conv(self.string), ignore_white_space=True):
            L.append((index, word))
        self.assertEqual(L, self.correct_positons)


    def test_iter2(self):
        self.add_words_and_make_automaton()
        A = self.A
        self.assertEqual(A.kind, ahocorasick.AHOCORASICK)

        L = []
        for index, word in A.iter(conv(self.string), ignore_white_space=True, start=12):
            L.append((index, word))
        self.assertEqual(L, self.correct_positons_start_12)


    def test_wrong_keyword(self):
        self.add_words_and_make_automaton()
        A = self.A
        self.assertEqual(A.kind, ahocorasick.AHOCORASICK)

        with self.assertRaises(TypeError):
            A.iter(conv(self.string), ignore_white_space2=True)


class TestAutomatonIterInvalidate(TestAutomatonBase):
    "Test if searching iterator is invalidated when trie/automaton change"

    def test_iter1(self):
        A = self.add_words_and_make_automaton()

        it = A.iter(conv(self.string))
        w = next(it)
        A.add_word(conv("should fail"), 1)
        with self.assertRaises(ValueError):
            w = next(it)


    def test_iter2(self):
        A = self.add_words_and_make_automaton()

        it = A.iter(conv(self.string))
        w = next(it)
        A.clear()
        with self.assertRaises(ValueError):
            w = next(it)


print_dumps = False

class TestPickle(TestAutomatonBase):
    "Test pickling/unpickling"

    def test_pickle(self):
        import pickle

        A = self.add_words_and_make_automaton();
        reduced = A.__reduce__()
        self.assertEqual(len(reduced), 2)
        if print_dumps:
            print(pickle.dumps(A))


    def test_unpickle(self):
        import pickle
        A = self.add_words_and_make_automaton();
        dump = pickle.dumps(A)
        B = pickle.loads(dump)

        self.compare_automatons(A, B)


    def test_unicode(self):
        # sample Russian words from issue #8
        import pickle

        test_sentences_rus = ["!ASM Print",
        "!ASM Print, tyre компания er",
        "!ASM Print, рекламно-производственная компания rr",
        "!Action Pact!",
        "!T.O.O.H.!",
        "!YES, лингвистический центр",
        "!ts, магазин",
        "!ФЕСТ",
        '"100-th" department store',
        '"1000 мелочей"',
        '"1001 мелочь"',
        '"19 отряд Федеральной противопожарной службы по Ленинградской области"',
        '"У Друзей"',
        '"ШТОРЫ и не только..."']

        A = ahocorasick.Automaton()
        for sentences in test_sentences_rus[-7:]:
            for index, word in enumerate(sentences.split(' ')):
                A.add_word(word, (index, word))

        dump = pickle.dumps(A)
        B = pickle.loads(dump)

        self.compare_automatons(A, B)


    def test_empty(self):
        import pickle

        A = ahocorasick.Automaton()
        dump = pickle.dumps(A)
        B = pickle.loads(dump)

        self.compare_automatons(A, B)


    def compare_automatons(self, A, B):
        if print_dumps:
            print([x for x in B.items()])
            print([x for x in A.items()])

        self.assertEqual(len(A), len(B))

        for item in zip(A.items(), B.items()):
            (AK, AV), (BK, BV) = item

            self.assertEqual(AK, BK)
            self.assertEqual(AV, BV)


class TestPickleStoreInts(TestCase):
    "Test pickling/unpickling for automaton of kind STORE_INTS/STORE_LEN"


    def add_words_and_make_automaton(self):
        A = ahocorasick.Automaton(ahocorasick.STORE_INTS)
        words = "tree trie bark branch barrier brag".split()

        for index, word in enumerate(words):
            A.add_word(word, index)

        A.make_automaton()

        return A


    def test_pickle_and_unpickle(self):
        import pickle
        A = self.add_words_and_make_automaton();
        dump = pickle.dumps(A)
        B = pickle.loads(dump)

        self.compare_automatons(A, B)


    def compare_automatons(self, A, B):
        if print_dumps:
            print([x for x in B.items()])
            print([x for x in A.items()])

        self.assertEqual(len(A), len(B))

        for item in zip(A.items(), B.items()):
            (AK, AV), (BK, BV) = item

            self.assertEqual(AK, BK)
            self.assertEqual(AV, BV)


class TestTrieStoreInts(TestCase):
    "Test storing plain ints as values (instead of python objects)"

    def setUp(self):
        self.A = ahocorasick.Automaton(ahocorasick.STORE_INTS);
        self.words = "word python aho corasick \x00\x00\x00".split()


    def test_add_word1(self):
        A = self.A

        # by default next values are stored
        for word in self.words:
            A.add_word(conv(word))

        I = list(range(1, len(self.words) + 1))
        L = [A.get(conv(word)) for word in self.words]
        self.assertEqual(I, L)


    def test_add_word2(self):
        A = self.A

        # store arbitrary ints
        for i, word in enumerate(self.words):
            A.add_word(conv(word), i + 123)

        I = list(range(123, 123 + len(self.words)))
        L = [A.get(conv(word)) for word in self.words]
        self.assertEqual(I, L)


    def test_add_word3(self):
        # not a number
        with self.assertRaises(TypeError):
            self.A.add_word(conv("xyz"), None)


    def test_iter(self):
        A = self.A
        for word in self.words:
            A.add_word(conv(word));

        I = set(range(1, len(A) + 1))
        L1 = [val for val in A.values()]
        L2 = [val for key, val in A.items()]

        self.assertEqual(L1, L2)
        self.assertEqual(set(L1), I)


    def test_find_all_and_iter(self):
        words = "he her hers she".split()
        string = "_sherhershe_"

        A = self.A
        for word in words:
            A.add_word(conv(word))

        A.make_automaton()

        # find_all()
        C = []
        def callback(index, value):
            C.append((index, value))

        A.find_all(conv(string), callback);

        # iter()
        L = [(index, value) for index, value in A.iter(conv(string))]

        #
        self.assertEqual(C, L)


class TestTrieStoreLengths(TestCase):
    """Test storing plain ints -- length of words --- as values
    (instead of python objects)"""

    def setUp(self):
        self.A = ahocorasick.Automaton(ahocorasick.STORE_LENGTH);
        self.words = "word python aho corasick \x00\x00\x00".split()


    def test_add_word1(self):
        A = self.A

        # by default next values are stored
        for word in self.words:
            A.add_word(conv(word))

        for key, value in A.items():
            self.assertEqual(len(key), value)


class TestSizeOf(TestCase):
    def setUp(self):
        self.A = ahocorasick.Automaton();
        words = "word python aho corasick tree bark branch root".split()
        for word in words:
            self.A.add_word(conv(word), 1)


    def test_sizeof(self):
        size1 = sys.getsizeof(self.A)

        # grow memory
        self.A.add_word("kitten", "fluffy")

        size2 = sys.getsizeof(self.A)

        # just change the assigned value, no changes to the trie structure
        self.A.add_word("word", "other value")

        size3 = sys.getsizeof(self.A)

        self.assertTrue(size2 > size1)
        self.assertTrue(size3 == size2)


class TestBugAutomatonSearch(TestAutomatonBase):
    """Bug in search"""

    def setUp(self):
        self.A = ahocorasick.Automaton()
        self.words = ['GT-C3303', 'SAMSUNG-GT-C3303K/']


    def test_bug(self):
        self.add_words_and_make_automaton()
        text = 'SAMSUNG-GT-C3303i/1.0 NetFront/3.5 Profile/MIDP-2.0 Configuration/CLDC-1.1'

        res = list(self.A.iter(conv(text)))

        self.assertEqual([(15, 'GT-C3303')], res)


class TestIntSequenceBase(TestCase):
    def setUp(self):
        self.A = ahocorasick.Automaton(ahocorasick.STORE_ANY, ahocorasick.KEY_SEQUENCE);


class TestIntSequence__TrieMethods(TestIntSequenceBase):

    def test_add__case_1(self):
        A = self.A

        ret = A.add_word((1, 2, 3), "foo")
        self.assertTrue(ret)
        self.assertTrue(A.kind == ahocorasick.TRIE)

        self.assertEqual(len(A), 1)
        self.assertTrue((1, 2, 3) in A)


    def test_add__case_2(self):
        A = self.A

        A.add_word((1, 2, 3), "foo")
        ret = A.add_word((1, 2, 3), "bar")
        self.assertFalse(ret)


    def test_add__case_3(self):
        A = self.A

        A.add_word((1, 2, 3), "foo")
        A.add_word((1, 2, 3, 4, 5), "bar")
        A.add_word((1, 3, 4, 5), "baz")

        self.assertEqual(len(A), 3);
        self.assertEqual(A.get((1, 2, 3)),          "foo");
        self.assertEqual(A.get((1, 2, 3, 4, 5)),    "bar");
        self.assertEqual(A.get((1, 3, 4, 5)),       "baz");


    def test_add__case_4(self):
        A = self.A

        ret = A.add_word((), "foo")
        self.assertFalse(ret)


    def test_add__case_5__wrong_argument_type(self):
        A = self.A

        with self.assertRaises(TypeError) as e:
            A.add_word("hello!", "foo")

        self.assertEqual(str(e.exception), "argument is not a supported sequence type")


    def test_add__case_6__wrong_item_type(self):
        A = self.A

        with self.assertRaises(ValueError) as e:
            A.add_word((1, 2, "hello!"), "foo")

        self.assertEqual(str(e.exception), "item #2 is not a number")


    def test_add__case_7__wrong_value(self):
        A = self.A

        with self.assertRaises(ValueError) as e:
            A.add_word((1, -1, 12), "foo")

        errmsg = str(e.exception)
        msgs = [
            "item #1: value -1 outside range [0..65535]",
            "item #1: value -1 outside range [0..4294967295]",
        ]

        self.assertIn(errmsg, msgs)


    def test_add__case_8__wrong_value(self):
        A = self.A

        with self.assertRaises(ValueError) as e:
            A.add_word((2**42, 0, 12), "foo")

        # Depending on python's version the message might be different,
        # but the type remains the same.

        errmsg = str(e.exception)
        msgs = [
            "item #0: value 4398046511104 outside range [0..65535]",
            "item #0: value 4398046511104 outside range [0..4294967295]",
            "item #0 is not a number",
        ]

        self.assertIn(errmsg, msgs)


    def test_match(self):
        A = self.A

        ret = A.add_word((1, 2, 3), "foo")
        self.assertTrue(A.match((1,)))
        self.assertTrue(A.match((1, 2)))
        self.assertTrue(A.match((1, 2, 3)))


    def test_longest_prefix(self):
        A = self.A

        ret = A.add_word((1, 2, 3, 4, 5, 6), "foo")
        self.assertEqual(A.longest_prefix((1, 2, 3, 111, 1111, 11111)), 3);
        self.assertEqual(A.longest_prefix((111, 1111, 11111)), 0);

    def test_iter1(self):
        A = self.A

        A.add_word((1, 2, 3), "foo")
        A.add_word((2, 3, 4, 5), "bar")
        A.add_word((2, 3, 5), "baz")
        A.make_automaton()

        L = [(index, value) for index, value in A.iter((1, 2, 3, 5))]

        self.assertEqual(L, [
            (2, "foo"),
            (3, "baz"),
        ])

    def test_iter2(self):
        A = self.A

        A.add_word((43, 89), (43, 89))
        A.add_word((43, 89, 64), (43, 89, 64))
        A.add_word((89, 64), (89, 64))
        A.add_word((89, 100), (89, 100))
        A.make_automaton()

        L = [
            (index, value)
            for index, value in
            A.iter((80, 80, 43, 89, 90, 89, 64, 100, 43, 89, 100))
        ]

        self.assertEqual(L, [
            (3, (43, 89)),
            (6, (89, 64)),
            (9, (43, 89)),
            (10, (89, 100)),
        ])


    def test_iter_wrong_argument_type(self):
        A = self.A
        A.add_word((89, 100), (89, 100))
        A.make_automaton()

        with self.assertRaisesRegex(TypeError, "tuple required"):
            self.A.iter(None)


class TestDump(TestAutomatonBase):
    def test_dump_empty(self):
        self.assertIsNone(self.A.dump())


    def test_dump_trie(self):
        self.add_words()
        ret = self.A.dump()

        self.assertEqual(3, len(ret))
        self.assertNotEmpty(ret[0])     # list of nodes
        self.assertNotEmpty(ret[1])     # list of edges
        self.assertEmpty(ret[2])        # list of fail links -- empty, if not an automaton


    def test_dump_automaton(self):
        self.add_words_and_make_automaton()
        ret = self.A.dump()

        self.assertEqual(3, len(ret))
        self.assertNotEmpty(ret[0])     # list of nodes
        self.assertNotEmpty(ret[1])     # list of edges
        self.assertNotEmpty(ret[2])     # list of fail links


class TestIssue53(TestCase):
    """
    Problems with handling of UCS-2 encoding
    """

    def test_case1(self):
        # test contributed by @woakesd (David Woakes)

        a = ahocorasick.Automaton()
        a.add_word('test', 'test')

        a.make_automaton()

        test_string = 'test 🙈 test?!'

        # wrongly calculated matching position
        for item in a.iter(test_string):
            start = item[0] - len(item[1]) + 1
            match = test_string[start:item[0] + 1]
            self.assertEqual(match, "test")


    def test_case2(self):
        a = ahocorasick.Automaton()
        a.add_word('test', 'test')

        a.make_automaton()

        test_string = '🙈' * 1000

        # wrongly calculated the input's length
        for item in a.iter(test_string):
            pass


class TestIssue68(TestCase):
    """
    Test problems with pickling
    """

    def test_case1(self):
        if _pickle is None:
            print("module _pickle not available")
            return

        A = ahocorasick.Automaton()
        for i in range(0, 65):
            A.add_word(str(i), (i, i))

        path = 'TestIssue68.test_case1'
        with open(path, 'wb') as f:
            _pickle.dump(A, f)

        with open(path, 'rb') as f:
            _pickle.load(f)

        try:
            os.unlink(path)
        except:
            pass


class TestLoadSave(TestAutomatonBase):

    def __init__(self, *args):
        super(TestAutomatonBase, self).__init__(*args)

        if os.path.isdir("/dev/shm"):
            tmp = "/dev/shm"
        else:
            tmp = "/tmp"

        self.path = conv(os.path.join(tmp, "test.dat"))


    def test_save__invalid_number_of_arguments(self):
        A = self.add_words_and_make_automaton();
        with self.assertRaisesRegex(ValueError, "expected exactly two arguments"):
            A.save()


    def test_save__invalid_argument_1(self):
        A = self.add_words_and_make_automaton();
        with self.assertRaisesRegex(TypeError, "the first argument must be a string"):
            A.save(None, None)


    def test_save__invalid_argument_2(self):
        A = self.add_words_and_make_automaton();
        with self.assertRaisesRegex(TypeError, "the second argument must be a callable object"):
            A.save(self.path, None)


    def test_load__invalid_number_of_arguments(self):
        with self.assertRaisesRegex(ValueError, "expected exactly two arguments"):
            ahocorasick.load()


    def test_load__invalid_argument_1(self):
        with self.assertRaisesRegex(TypeError, "the first argument must be a string"):
            ahocorasick.load(None, None)


    def test_load__invalid_argument_2(self):
        with self.assertRaisesRegex(TypeError, "the second argument must be a callable object"):
            ahocorasick.load("/dev/shm/test.dump", None)


    def test_save(self):
        import pickle

        A = self.add_words_and_make_automaton();

        A.save(self.path, pickle.dumps)


    def test_save_and_load_empty(self):
        import pickle

        A = ahocorasick.Automaton()

        A.save(self.path, pickle.dumps)
        B = ahocorasick.load(self.path, pickle.loads)

        self.compare_automatons(A, B)


    def test_save_and_load_trie(self):
        import pickle

        A = self.add_words()

        A.save(self.path, pickle.dumps)
        B = ahocorasick.load(self.path, pickle.loads)

        self.compare_automatons(A, B)


    def test_save_and_load_automaton(self):
        import pickle

        A = self.add_words_and_make_automaton();

        A.save(self.path, pickle.dumps)
        B = ahocorasick.load(self.path, pickle.loads)

        self.compare_automatons(A, B)


    def test_save_ints(self):
        A = ahocorasick.Automaton(ahocorasick.STORE_INTS)
        with self.assertRaisesRegex(ValueError, "expected exactly one argument"):
            A.save(self.path, None)


    def test_save_and_load_ints(self):
        import pickle

        A = ahocorasick.Automaton(ahocorasick.STORE_INTS)
        for i, word in enumerate(conv("he she her cat car carriage zoo")):
            A.add_word(word, i)

        A.save(self.path)
        B = ahocorasick.load(self.path, pickle.loads)

        self.compare_automatons(A, B)


    def compare_automatons(self, A, B):
        if print_dumps:
            print([x for x in B.items()])
            print([x for x in A.items()])

        self.assertEqual(len(A), len(B))

        A = list(A.items())
        B = list(B.items())

        for item in zip(A, B):
            (AK, AV), (BK, BV) = item

            self.assertEqual(AK, BK)
            self.assertEqual(AV, BV)


if __name__ == '__main__':
    unittest.main()