internals_test.py

from internals import format_assistant_reply, format_openai_message_content


def test_format_assistant_reply():
    for (content, expected) in [
        (
            "\n\nSorry, I cannot answer the question.",
            "Sorry, I cannot answer the question.",
        ),
        ("\n\n```python\necho 'foo'\n```", "```\necho 'foo'\n```"),
        ("\n\n```ruby\nputs 'foo'\n```", "```\nputs 'foo'\n```"),
        (
            "\n\n```java\nSystem.out.println(123);\n```",
            "```\nSystem.out.println(123);\n```",
        ),
    ]:
        result = format_assistant_reply(content)
        assert result == expected


def test_format_openai_message_content():
    # https://github.com/seratch/ChatGPT-in-Slack/pull/5
    for (content, expected) in [
        (
            """#include <stdio.h>
int main(int argc, char *argv[])
{
    printf("Hello, world!\n");
    return 0;
}""",
            """#include <stdio.h>
int main(int argc, char *argv[])
{
    printf("Hello, world!\n");
    return 0;
}""",
        ),
    ]:
        result = format_openai_message_content(content)
        assert result == expected