diff --git a/src/main/kotlin/net/woggioni/klevtree/LevTrie.kt b/src/main/kotlin/net/woggioni/klevtree/LevTrie.kt
index 1a3aab1..b898b38 100644
--- a/src/main/kotlin/net/woggioni/klevtree/LevTrie.kt
+++ b/src/main/kotlin/net/woggioni/klevtree/LevTrie.kt
@@ -102,46 +102,39 @@ interface ILevTrie : ICharTrie<IntArray> {
     fun fuzzySearch(word : String, maxResult: Int) : List<Pair<String, Int>> {
         val result = sortedSetOf<Pair<String, Int>>(compareBy({ it.second }, { it.first }))
         val requiredSize = word.length + 1
-        fun visitNode(stack: List<StackContext<LevNode, Unit>>) : TreeNodeVisitor.VisitOutcome {
-            if(stack.size > 1) {
+        val visitor = object: TreeNodeVisitor<LevNode, Unit> {
+            override fun visitPre(stack: List<StackContext<LevNode, Unit>>): TreeNodeVisitor.VisitOutcome {
                 val currentStackElement = stack.last()
-                if(currentStackElement.node.key == null) {
-                    val sb = StringBuilder()
-                    for(c in currentStackElement.node.linealDescendant()) {
-                        sb.append(c)
-                    }
-                    val candidate = sb.toString()
-                    val distance = stack[stack.size - 2].node.payload!![word.length]
-                    result.add(candidate to distance)
-                    if(result.size > maxResult) {
-                        result.remove(result.last())
-                    }
-                    return TreeNodeVisitor.VisitOutcome.SKIP
-                } else {
-                    return distanceCalculator.compute(keyChecker, stack, word,
-                        if(result.size == maxResult) result.last().second else -1)
-                }
-            } else {
-                return TreeNodeVisitor.VisitOutcome.CONTINUE
-            }
-        }
-        val visitor = if(root.payload == null || root.payload!!.size < requiredSize) {
-            object: TreeNodeVisitor<LevNode, Unit> {
-                override fun visitPre(stack: List<StackContext<LevNode, Unit>>): TreeNodeVisitor.VisitOutcome {
-                    val currentNode = stack.last()
+                val currentNode = currentStackElement.node
+                if(currentNode.payload == null ||
+                    currentNode.payload!!.size < requiredSize) {
                     if(stack.size == 1) {
-                        currentNode.node.payload = IntArray(requiredSize) { i -> i }
+                        currentNode.payload = IntArray(requiredSize) { i -> i }
                     } else {
-                        currentNode.node.payload = IntArray(requiredSize) { i -> if(i == 0) stack.size - 1 else 0 }
+                        currentNode.payload = IntArray(requiredSize) { i -> if(i == 0) stack.size - 1 else 0 }
                     }
-                    visitNode(stack)
+                }
+                if(stack.size > 1) {
+                    if(currentStackElement.node.key == null) {
+                        val sb = StringBuilder()
+                        for(c in currentStackElement.node.linealDescendant()) {
+                            sb.append(c)
+                        }
+                        val candidate = sb.toString()
+                        val distance = stack[stack.size - 2].node.payload!![word.length]
+                        result.add(candidate to distance)
+                        if(result.size > maxResult) {
+                            result.remove(result.last())
+                        }
+                        return TreeNodeVisitor.VisitOutcome.SKIP
+                    } else {
+                        return distanceCalculator.compute(keyChecker, stack, word,
+                            if(result.size == maxResult) result.last().second else -1)
+                    }
+                } else {
                     return TreeNodeVisitor.VisitOutcome.CONTINUE
                 }
             }
-        } else object: TreeNodeVisitor<LevNode, Unit> {
-            override fun visitPre(stack: List<StackContext<LevNode, Unit>>): TreeNodeVisitor.VisitOutcome {
-                return visitNode(stack)
-            }
         }
         val walker = TreeWalker<LevNode, Unit>(visitor)
         walker.walk(root)
diff --git a/src/test/kotlin/net/woggioni/klevtree/LevtreeTest.kt b/src/test/kotlin/net/woggioni/klevtree/LevtreeTest.kt
index 27d1979..a30ee2f 100644
--- a/src/test/kotlin/net/woggioni/klevtree/LevtreeTest.kt
+++ b/src/test/kotlin/net/woggioni/klevtree/LevtreeTest.kt
@@ -1,33 +1,111 @@
 package net.woggioni.klevtree
 
+import org.junit.Assert
 import org.junit.Test
-import java.io.BufferedReader
-import java.io.InputStreamReader
 
 
 class LevtreeTest {
 
     @Test
-    fun foo() {
-        val reader = BufferedReader(
-            InputStreamReader(javaClass.getResourceAsStream("/cracklib-small")))
+    fun trivialTest() {
         val tree = LevTrie()
         tree.caseSensitive = false
-        try {
-            for(line in reader.lines()) {
-                tree.add(line.asIterable())
-            }
-        } finally {
-            reader.close()
-        }
-        println(tree.add("dailies"))
-        var node = tree.search("dailies")
-        println(node!!.linealDescendant())
-        tree.remove("dailies")
-        node = tree.search("dailies")
-        println(node)
         tree.algorithm = ILevTrie.Algorithm.DAMERAU_LEVENSHTEIN
-        val result = tree.fuzzySearch("daiiles", 5)
-        println(result)
+        val word = "dailies"
+        run {
+            val pair = tree.add(word)
+            Assert.assertTrue(pair.first)
+            val node = tree.search(word)
+            Assert.assertNotNull(node)
+            Assert.assertEquals(
+                word,
+                node!!.linealDescendant().fold(StringBuilder(), StringBuilder::append).toString()
+            )
+            val result = tree.fuzzySearch(word, 5)
+            Assert.assertEquals(1, result.size)
+            Assert.assertEquals(word to 0, result[0])
+        }
+        run {
+            tree.remove(word)
+            val node = tree.search(word)
+            Assert.assertNull(node)
+            val result = tree.fuzzySearch(word, 5)
+            Assert.assertEquals(0, result.size)
+        }
+    }
+
+
+    private fun initLevTrie() : LevTrie {
+        val words = listOf(
+            "tired",
+            "authorise",
+            "exercise",
+            "bloody",
+            "ritual",
+            "trail",
+            "resort",
+            "landowner",
+            "navy",
+            "captivate",
+            "captivity",
+            "north")
+        return run {
+            val res = LevTrie()
+            words.forEach {res.add(it)}
+            res
+        }
+    }
+
+    @Test
+    fun levenshteinDistanceTest() {
+        val tree = initLevTrie()
+        tree.caseSensitive = false
+        tree.algorithm = ILevTrie.Algorithm.LEVENSHTEIN
+        run {
+            val word = "fired"
+            val result = tree.fuzzySearch(word, 4)
+            Assert.assertEquals(4, result.size)
+            Assert.assertEquals("tired" to 1, result[0])
+        }
+        run {
+            val word = "tierd"
+            val result = tree.fuzzySearch(word, 4)
+            Assert.assertEquals(4, result.size)
+            Assert.assertEquals("tired" to 2, result[0])
+        }
+        run {
+            val word = "tierd"
+            tree.remove("tired")
+            val result = tree.fuzzySearch(word, 4)
+            Assert.assertEquals(4, result.size)
+            Assert.assertEquals("trail" to 4, result[0])
+        }
+    }
+
+    @Test
+    fun damerauLevenshteinDistanceTest() {
+        val tree = initLevTrie()
+        tree.caseSensitive = false
+        tree.algorithm = ILevTrie.Algorithm.DAMERAU_LEVENSHTEIN
+        run {
+            val word = "fired"
+            val result = tree.fuzzySearch(word, 4)
+            Assert.assertEquals(4, result.size)
+            Assert.assertEquals("tired" to 1, result[0])
+        }
+        run {
+            val word = "capitvate"
+            val result = tree.fuzzySearch(word, 4)
+            Assert.assertEquals(4, result.size)
+            Assert.assertEquals("captivate" to 1, result[0])
+            Assert.assertEquals("captivity" to 3, result[1])
+        }
+        run {
+            tree.remove("captivate")
+            val word = "capitvate"
+            val result = tree.fuzzySearch(word, 4)
+            Assert.assertEquals(4, result.size)
+            Assert.assertEquals("captivity" to 3, result[0])
+        }
     }
 }
\ No newline at end of file