parse chat messages tests

1 files changed, 140 insertions, 0 deletions

diff --git a/tests/chat_test.py b/tests/chat_test.py
new file mode 100644
index 0000000..9acfecb
--- /dev/null
+++ b/tests/chat_test.py
@@ -0,0 +1,140 @@
+from utils import parse_chat_messages
+import os
+
+curr_dir = os.path.dirname(__file__)
+
+def strip_text(txt):
+    txt = txt.strip()
+    lines = txt.splitlines()
+    return "\n".join([line.lstrip() for line in lines])
+
+def test_parse_user_message():
+    chat_content = strip_text(
+    """
+    >>> user
+
+    generate lorem ipsum
+    """)
+    messages = parse_chat_messages(chat_content)
+    assert 1 == len(messages)
+    assert 'user' == messages[0]['role']
+    assert 'generate lorem ipsum' == messages[0]['content']
+
+
+def test_parse_system_message():
+    chat_content = strip_text("""
+    >>> system
+
+    you are general assystant
+
+    >>> user
+
+    generate lorem ipsum
+    """)
+    messages = parse_chat_messages(chat_content)
+    assert 2 == len(messages)
+    assert 'system' == messages[0]['role']
+    assert 'you are general assystant' == messages[0]['content']
+    assert 'user' == messages[1]['role']
+    assert 'generate lorem ipsum' == messages[1]['content']
+
+def test_parse_assistant_message():
+    chat_content = strip_text("""
+    >>> user
+
+    generate lorem ipsum
+
+    <<< assistant
+
+    bla bla bla
+
+    >>> user
+
+    again
+    """)
+    messages = parse_chat_messages(chat_content)
+    assert 3 == len(messages)
+    assert 'user' == messages[0]['role']
+    assert 'generate lorem ipsum' == messages[0]['content']
+    assert 'assistant' == messages[1]['role']
+    assert 'bla bla bla' == messages[1]['content']
+    assert 'user' == messages[2]['role']
+    assert 'again' == messages[2]['content']
+
+def test_parse_include_single_file_message():
+    chat_content = strip_text(f"""
+    >>> user
+
+    translate to human language
+
+    >>> include
+
+    {curr_dir}/resources/test1.include.txt
+    """)
+    messages = parse_chat_messages(chat_content)
+    assert 2 == len(messages)
+    assert 'user' == messages[0]['role']
+    assert 'translate to human language' == messages[0]['content']
+    assert 'user' == messages[1]['role']
+    expected_content = strip_text(f"""
+    ==> {curr_dir}/resources/test1.include.txt <==
+    hello world
+    """)
+    assert expected_content == messages[1]['content']
+
+def test_parse_include_multiple_files_message():
+    chat_content = strip_text(f"""
+    >>> user
+
+    translate to human language
+
+    >>> include
+
+    {curr_dir}/resources/test1.include.txt
+    {curr_dir}/resources/test2.include.txt
+    """)
+    messages = parse_chat_messages(chat_content)
+    assert 2 == len(messages)
+    assert 'user' == messages[0]['role']
+    assert 'translate to human language' == messages[0]['content']
+    assert 'user' == messages[1]['role']
+    expected_content = strip_text(f"""
+    ==> {curr_dir}/resources/test1.include.txt <==
+    hello world
+
+    ==> {curr_dir}/resources/test2.include.txt <==
+    vim is awesome
+    """)
+    assert expected_content == messages[1]['content']
+
+def test_parse_include_glob_files_message():
+    chat_content = strip_text(f"""
+    >>> user
+
+    translate to human language
+
+    >>> include
+
+    {curr_dir}/**/*.include.txt
+    """)
+    messages = parse_chat_messages(chat_content)
+    assert 2 == len(messages)
+    assert 'user' == messages[0]['role']
+    assert 'translate to human language' == messages[0]['content']
+    assert 'user' == messages[1]['role']
+    expected_content = strip_text(f"""
+    ==> {curr_dir}/resources/test1.include.txt <==
+    hello world
+
+    ==> {curr_dir}/resources/test2.include.txt <==
+    vim is awesome
+    """)
+    assert expected_content == messages[1]['content']
+
+def test_parse_include_image_message():
+    # TODO
+    pass
+
+def test_parse_include_image_with_files_message():
+    # TODO
+    pass