Generate search tree tests

Create unit tests with edge cases for binary search trees

You can query a model directly and test the results returned when using different parameter values with the Cloud console, or by calling the Vertex AI API directly.

System instructions

You are an experienced programmer, well-versed in various common algorithms. Your task is to create sufficiently detailed unit tests for the implementation of a binary search tree to ensure the correctness of this class's implementation. Do not offer explanations, just create the unit tests.

Prompt Freeform

I have a code implementation of a binary search tree, but I'm not sure if it's correct. To confirm that, I hope you can write unit tests for me to test each function (except for private functions) in this class. Your unit tests must include edge cases such as performing operations on an empty tree and handling duplicate values. Additionally, your code should contain a main function.

bst.py

from dataclasses import dataclass
from typing import Optional

@dataclass
class TreeNode:
    value: int
    left: Optional['TreeNode'] = None
    right: Optional['TreeNode'] = None

class BST:
    def __init__(self):
        self.root: Optional[TreeNode] = None

    def insert(self, value: int) -> None:
        if not self.root:
            self.root = TreeNode(value)
        else:
            self._insert_helper(self.root, value)

    def _insert_helper(self, node: TreeNode, value: int) -> None:
        if value < node.value:
            if node.left is None:
                node.left = TreeNode(value)
            else:
                self._insert_helper(node.left, value)
        else:
            if node.right is None:
                node.right = TreeNode(value)
            else:
                self._insert_helper(node.right, value)

    def find(self, value: int) -> Optional[TreeNode]:
        return self._find_helper(self.root, value)

    def _find_helper(self, node: Optional[TreeNode], value: int) -> Optional[TreeNode]:
        if node is None or node.value == value:
            return node
        if value < node.value:
            return self._find_helper(node.left, value)
        return self._find_helper(node.right, value)

    def _lookup_min(self, node: TreeNode) -> TreeNode:
        current = node
        while current.left:
            current = current.left
        return current

    def print_inorder(self) -> None:
        self._print_inorder_helper(self.root)
        print()

    def _print_inorder_helper(self, node: Optional[TreeNode]) -> None:
        if node:
            self._print_inorder_helper(node.left)
            print(node.value, end=' ')
            self._print_inorder_helper(node.right)

Response

import unittest
from bst import BST, TreeNode

class TestBST(unittest.TestCase):

    def test_insert_empty(self):
        bst = BST()
        bst.insert(5)
        self.assertEqual(bst.root.value, 5)

    def test_insert_multiple(self):
        bst = BST()
        bst.insert(5)
        bst.insert(3)
        bst.insert(7)
        bst.insert(1)
        bst.insert(4)
        bst.insert(6)
        bst.insert(8)
        self.assertEqual(bst.root.value, 5)
        self.assertEqual(bst.root.left.value, 3)
        self.assertEqual(bst.root.right.value, 7)

    def test_insert_duplicates(self):
        bst = BST()
        bst.insert(5)
        bst.insert(5)
        self.assertEqual(bst.root.value, 5)
        self.assertIsNone(bst.root.left)
        self.assertIsNone(bst.root.right)

    def test_find_empty(self):
        bst = BST()
        self.assertIsNone(bst.find(5))

    def test_find_present(self):
        bst = BST()
        bst.insert(5)
        bst.insert(3)
        bst.insert(7)
        self.assertEqual(bst.find(5).value, 5)
        self.assertEqual(bst.find(3).value, 3)
        self.assertEqual(bst.find(7).value, 7)

    def test_find_absent(self):
        bst = BST()
        bst.insert(5)
        bst.insert(3)
        bst.insert(7)
        self.assertIsNone(bst.find(1))
        self.assertIsNone(bst.find(9))

    def test_find_duplicates(self):
        bst = BST()
        bst.insert(5)
        bst.insert(5)
        self.assertEqual(bst.find(5).value, 5)

    def test_print_inorder_empty(self):
        bst = BST()
        bst.print_inorder()
        
    def test_print_inorder_non_empty(self):
        bst = BST()
        bst.insert(5)
        bst.insert(3)
        bst.insert(7)
        bst.insert(1)
        bst.insert(4)
        bst.insert(6)
        bst.insert(8)
        bst.print_inorder()


if __name__ == '__main__':
    unittest.main()
Model: gemini-1.5-flash-002
Temperature: 0.2
Max output tokens: 8192
TopK: 40.0
TopP: 0.95