added BST, but remove not working

recursion/recursive_datastructure/solution/test_search_trees.py

@@ -1,24 +1,51 @@
-from search_trees import BinarySearchTree
+from typing import Iterator, Optional
+from search_trees import BinarySearchTree, insert, remove, exists
+from trees import Node
 from random import randint as random
 
-MAX= 1000
+
+def traverse[T](tree: Optional[Node[T]]) -> Iterator[T]:
+    match tree:
+        case Node(value, left, right):
+            yield from traverse(left)
+            yield value
+            yield from traverse(right)
+        case _:
+            return
+
+
+MAX = 100
 MIN = 0
 LENGTH = 100
 
+NUMS = {random(MIN, MAX) for _ in range(0, LENGTH)}
+NUMS_FILTER = {random(MIN, MAX) for _ in range(0, LENGTH // 4)}
+
 def test_insert():
-    nums = [random(MIN, MAX) for _ in range(0, LENGTH)]
-    bst = BinarySearchTree()
-    for num in nums:
-        bst.insert(num)
-    assert list(iter(bst)) == sorted(nums)
-    
+    bst: BinarySearchTree = None
+    for num in NUMS:
+        bst = insert(bst, num)
+    assert list(traverse(bst)) == sorted(NUMS)
+
+
 def test_remove():
-    nums = {random(MIN, MAX) for _ in range(0, LENGTH)}
-    fil = {random(MIN, MAX) for _ in range(0, LENGTH // 4)}
-    bst = BinarySearchTree()
-    for num in nums:
-        bst.insert(num)
-    for num in fil:
-        bst.remove(num)
-        
-    assert list(iter(bst)) == sorted(filter(lambda x: x not in fil, nums))
+    bst: BinarySearchTree = None
+    for num in NUMS:
+        bst = insert(bst, num)
+    assert list(traverse(bst)) == sorted(NUMS)
+    for num in NUMS_FILTER:
+        remove(bst, num)
+
+    assert list(traverse(bst)) == sorted(filter(lambda x: x not in NUMS_FILTER, NUMS))
+
+
+def test_exists():
+    bst: BinarySearchTree = None
+    for num in NUMS:
+        bst = insert(bst, num)
+    for num in NUMS_FILTER:
+        remove(bst, num)
+
+    assert all(map(lambda x: exists(bst, x),
+               filter(lambda x: x not in NUMS_FILTER, NUMS)))
+    assert all(map(lambda x: not exists(bst, x), NUMS_FILTER))