Update test_data_structure.py

Author: antmarakis

tests/test_data_structure.py

@@ -7,7 +7,8 @@
     linked_list,
     tree,
     graph,
-    heap)
+    heap,
+    trie)
 
 
 class TestStack(unittest.TestCase):
@@ -237,5 +238,22 @@ def test_heap(self):
         self.assertEqual(myHeap.queue, expectedResult)
 
 
+class TestTrie(unittest.TestCase):
+    def test_stack(self):
+        myTrie = trie.Trie()
+        myTrie.Insert('the')
+        myTrie.Insert('turtle')
+        myTrie.Insert('thesaurus')
+        myTrie.Insert('chocolate')
+        myTrie.Insert('flying')
+
+        self.assertEqual(myTrie.FindWords('th'), ['the', 'thesaurus'])
+        self.assertEqual(myTrie.FindWords('e'), None)
+
+        self.assertEqual(myTrie.Search('chocolate'), True)
+        self.assertEqual(myTrie.Search('flying'), True)
+        self.assertEqual(myTrie.Search('walking'), False)
+
+
 if __name__ == '__main__':
     unittest.main()