def test_build_prompt():
# Test simple prompt
prompt = '[?] name: (Heisenberg) '
assert remove_ansi(build_prompt('name', '(Heisenberg)')) == prompt
# Test custom prompt
prompt = '[?] What is your name? (Heisenberg) '
built_prompt = build_prompt('What is your name?', '(Heisenberg)')
assert prompt == remove_ansi(built_prompt)
# Test highly customized prompt (default not in the end)
prompt = '[?] I [walk]ed to China '
built_prompt = build_prompt('I {default}ed to China', '[walk]')
assert prompt == remove_ansi(built_prompt)
def test_remove_ansi():
# Test a general case
ansi_string = '[\x1b[32m?\x1b[0m] name: '
assert remove_ansi(ansi_string) == '[?] name: '
